| """Calibration and a small eval for deg checkpoints (SPEC.md §11, §12). |
| |
| ``calibrate`` fits one temperature per head kind (choice, score, noul, extract) by NLL on |
| held-out data: datagen ``val`` records only by default (``--source datagen-val``; the |
| /v1/fill heads are calibrated for /v1/fill data). jev dev can be reported next to it |
| (``--report-jev``, fit pool unchanged) or fit on (``--source jev-dev`` or |
| ``datagen-val+jev-dev``, the run-1 pool). ``evaluate`` runs the full |
| :func:`deg.infer.fill` path on datagen ``val`` and ``test`` (held-out domains) records and |
| reports, per split: valid-response rate, decide accuracy + ECE, extract exact match |
| (positional, and ``extract_em_aligned`` with array items matched order-insensitively), array |
| count accuracy, and generated tokens vs the tokens of the gold JSON written out in full |
| (the full-JSON constrained-decoding baseline emits at least that many). jev dev decisions |
| are scored against the teacher argmax. |
| |
| python -m deg.evaluate --ckpt /data/deg/runs/smoke/final |
| python -m deg.evaluate calibrate --ckpt /data/deg/runs/run-1/step-0020000 \ |
| --out /data/deg/runs/run-1/final-v2 --report-jev |
| |
| ``calibrate`` writes a calibrated copy of a checkpoint (same backbone and head weights, the |
| fitted temperatures in ``heads.safetensors`` and ``deg.json``); the source is not modified. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import random |
| import shutil |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
|
|
| from deg.infer import fill |
| from deg.model import KINDS, DegModel, load_model |
| from deg.request import FillOptions |
| from deg.response import FillError |
| from deg.schema import compile_schema |
|
|
| TEMPS = [round(0.4 + 0.1 * i, 2) for i in range(37)] |
| CALIBRATION_SOURCES = ("datagen-val", "jev-dev", "datagen-val+jev-dev") |
| CALIBRATION_TARGETS = ("hard", "soft") |
| DECIDE_KINDS = ("choice", "score", "noul") |
|
|
|
|
| def _flatten(value: Any, prefix: str = "") -> dict[str, Any]: |
| if isinstance(value, dict): |
| out: dict[str, Any] = {} |
| for k, v in value.items(): |
| out |= _flatten(v, f"{prefix}/{k.replace('~', '~0').replace('/', '~1')}") |
| return out |
| if isinstance(value, list): |
| out = {f"{prefix}#len": len(value)} |
| for i, v in enumerate(value): |
| out |= _flatten(v, f"{prefix}/{i}") |
| return out |
| return {prefix: value} |
|
|
|
|
| def align_arrays(gold: Any, pred: Any, eq: Any = None) -> Any: |
| """``pred`` with its arrays reordered (recursively) so item i is the predicted item |
| that best matches gold item i (greedy by the fraction of matching non-null leaves, |
| ties by index); unmatched gold items become ``None`` and unmatched predicted items |
| are dropped. Order-insensitive scoring: anchored arrays come out in document order, |
| which need not be the gold's order.""" |
| eq = eq or (lambda g, p: g == p) |
| if isinstance(gold, dict) and isinstance(pred, dict): |
| return {k: (align_arrays(gold[k], v, eq) if k in gold else v) for k, v in pred.items()} |
| if isinstance(gold, list) and isinstance(pred, list): |
| pairs = [] |
| for i, g in enumerate(gold): |
| fg = {k: v for k, v in _flatten(g).items() if not k.endswith("#len") and v is not None} |
| for j, q in enumerate(pred): |
| fq = _flatten(q) |
| hit = sum(1 for k, v in fg.items() if k in fq and eq(v, fq[k])) |
| pairs.append((-hit / max(1, len(fg)), i, j)) |
| pairs.sort() |
| used_g: set[int] = set() |
| used_p: set[int] = set() |
| match: dict[int, int] = {} |
| for score, i, j in pairs: |
| if score >= 0 or i in used_g or j in used_p: |
| continue |
| used_g.add(i) |
| used_p.add(j) |
| match[i] = j |
| free = [j for j in range(len(pred)) if j not in used_p] |
| out = [] |
| for i, g in enumerate(gold): |
| if i in match: |
| out.append(align_arrays(g, pred[match[i]], eq)) |
| elif free: |
| out.append(align_arrays(g, pred[free.pop(0)], eq)) |
| return out + [pred[j] for j in free] |
| return pred |
|
|
|
|
| def _abstract(ptr: str) -> str: |
| return "/".join("*" if p.isdigit() else p for p in ptr.split("/")) |
|
|
|
|
| def ece(pairs: list[tuple[float, bool]], bins: int = 10) -> float | None: |
| if not pairs: |
| return None |
| total = 0.0 |
| for b in range(bins): |
| lo, hi = b / bins, (b + 1) / bins |
| sel = [(c, ok) for c, ok in pairs if lo < c <= hi or (b == 0 and c == 0)] |
| if sel: |
| conf = sum(c for c, _ in sel) / len(sel) |
| acc = sum(ok for _, ok in sel) / len(sel) |
| total += len(sel) / len(pairs) * abs(conf - acc) |
| return round(total, 4) |
|
|
|
|
| def _collect_logits(model: DegModel, examples: list[Any]) -> dict[str, list]: |
| """kind -> [(logits/start-end, target)] on the inference path.""" |
| got: dict[str, list] = {k: [] for k in KINDS} |
| for ex in examples: |
| items = [it for it in ex.items if it.query is not None] |
| if not items: |
| continue |
| hidden, state = model.trunk([ex.ctx]) |
| outs = model.answer_batched(state, hidden, ex.ctx, [it.query for it in items]) |
| for it, out in zip(items, outs): |
| if it.kind == "extract": |
| got["extract"].append(((out.start.float(), out.end.float()), it)) |
| elif it.kind in ("choice", "score", "noul"): |
| got[it.kind].append((out.logits.float(), it)) |
| return got |
|
|
|
|
| def _nll(kind: str, rows: list, t: float, hard: bool = False) -> float: |
| """Mean NLL at temperature ``t``; ``hard``: decide targets become one-hot of their |
| argmax (the filled value / the teacher's top outcome), the label accuracy and ECE are |
| scored against.""" |
| total = 0.0 |
| for logits, it in rows: |
| if kind == "extract": |
| start, end = logits |
| for lg, pos in ((start, it.starts), (end, it.ends)): |
| logp = torch.log_softmax(lg / t, -1) |
| idx = torch.tensor([lg.shape[0] - 1 if p < 0 else p for p in pos], |
| device=lg.device) |
| total -= float(torch.logsumexp(logp[idx.unique()], 0)) |
| else: |
| target = torch.tensor(it.target, device=logits.device).clamp_min(0) |
| if hard: |
| target = torch.nn.functional.one_hot(target.argmax(), len(target)).float() |
| total -= float((target / target.sum() * torch.log_softmax(logits / t, -1)).sum()) |
| return total / max(1, len(rows)) |
|
|
|
|
| def _top1_pairs(kind: str, rows: list, t: float) -> list[tuple[float, bool]]: |
| """(confidence, correct) of each decide row's argmax at temperature ``t``; correct = |
| the target's argmax (the filled value, or the teacher's top outcome).""" |
| out = [] |
| for logits, it in rows: |
| probs = torch.softmax(logits / t, -1) |
| best = max(range(len(it.target)), key=it.target.__getitem__) |
| out.append((float(probs.max()), int(probs.argmax()) == best)) |
| return out |
|
|
|
|
| def fit_report(got: dict[str, list], temps: dict[str, float], before: dict[str, float], |
| hard: bool = True) -> dict[str, Any]: |
| """Per head kind of one pool: rows, NLL and (decide kinds) top-1 accuracy + ECE at the |
| ``before`` and fitted temperatures, plus all decide kinds pooled.""" |
| out: dict[str, Any] = {} |
| pooled: dict[str, list] = {"before": [], "after": []} |
| for kind in KINDS: |
| rows = got.get(kind) or [] |
| if not rows: |
| continue |
| m: dict[str, Any] = {"n": len(rows), "t_before": before[kind], "t_after": temps[kind], |
| "nll_before": round(_nll(kind, rows, before[kind], hard), 4), |
| "nll_after": round(_nll(kind, rows, temps[kind], hard), 4)} |
| if kind in DECIDE_KINDS: |
| for tag, t in (("before", before[kind]), ("after", temps[kind])): |
| pairs = _top1_pairs(kind, rows, t) |
| pooled[tag] += pairs |
| m[f"ece_{tag}"] = ece(pairs) |
| m["acc"] = round(sum(ok for _, ok in pairs) / len(pairs), 4) |
| out[kind] = m |
| if pooled["after"]: |
| out["decide"] = {"n": len(pooled["after"]), "ece_before": ece(pooled["before"]), |
| "ece_after": ece(pooled["after"])} |
| return out |
|
|
|
|
| def _as_paths(records: Path | list[Path] | tuple[Path, ...]) -> list[Path]: |
| return [Path(records)] if isinstance(records, (str, Path)) else [Path(r) for r in records] |
|
|
|
|
| @torch.no_grad() |
| def calibrate(model: DegModel, records: Path | list[Path], jev_dev: Path, n_jev: int = 300, |
| seed: int = 0, info: dict[str, Any] | None = None, *, |
| source: str = "datagen-val", report_jev: bool = False, |
| max_context_tokens: int | None = None, target: str = "hard" |
| ) -> dict[str, float]: |
| """Fit and set ``model``'s temperatures on ``source``: datagen ``val`` records of every |
| ``records`` file (default), jev dev, or both pooled. ``info`` (if given) receives the |
| sample sizes and, per pool (``datagen_val``, ``jev_dev`` when fit on or |
| ``report_jev``), NLL/ECE before (the model's current temperatures) and after. |
| ``max_context_tokens`` leaves longer val contexts out (the end-of-training call). |
| ``target="hard"`` (default) fits decide heads to the one-hot filled value, the label |
| accuracy/ECE are judged against; ``"soft"`` fits the soft training targets (Clef |
| ``decide_probs``, jev teacher probs), which are less confident than the gold and pull |
| the temperatures above 1 (calibration-v2.md).""" |
| from deg.train import deg_example, jev_example, read_jsonl |
|
|
| if target not in CALIBRATION_TARGETS: |
| raise ValueError(f"calibration target {target!r} not in {CALIBRATION_TARGETS}") |
| hard = target == "hard" |
| if source not in CALIBRATION_SOURCES: |
| raise ValueError(f"calibration source {source!r} not in {CALIBRATION_SOURCES}") |
| fit_val = source in ("datagen-val", "datagen-val+jev-dev") |
| fit_jev = source in ("jev-dev", "datagen-val+jev-dev") |
| rng = random.Random(seed) |
| pools: dict[str, dict[str, list]] = {} |
| sizes: dict[str, int] = {} |
| if fit_val: |
| examples = [] |
| for path in _as_paths(records): |
| for r in read_jsonl(path, "val"): |
| ex = deg_example(model, r, rng, shuffle_p=0.0) |
| if max_context_tokens is None or len(ex.ctx.ids) <= max_context_tokens: |
| examples.append(ex) |
| sizes["val_records"] = len(examples) |
| pools["datagen_val"] = _collect_logits(model, examples) |
| if (fit_jev or report_jev) and jev_dev.exists() and n_jev: |
| examples = [e for r in read_jsonl(jev_dev, None, n_jev) |
| if (e := jev_example(model, r, rng, shuffle_p=0.0)) is not None] |
| sizes["jev_records"] = len(examples) |
| pools["jev_dev"] = _collect_logits(model, examples) |
| fit_on = [k for k, on in (("datagen_val", fit_val), ("jev_dev", fit_jev)) if on] |
| got = {k: [row for name in fit_on for row in pools.get(name, {}).get(k, [])] |
| for k in KINDS} |
| before = dict(zip(KINDS, (float(t) for t in model.heads.temperature.tolist()))) |
| temps = {} |
| for i, kind in enumerate(KINDS): |
| rows = got.get(kind) or [] |
| if len(rows) < 10: |
| temps[kind] = before[kind] |
| continue |
| temps[kind] = min(TEMPS, key=lambda t: _nll(kind, rows, t, hard)) |
| model.heads.temperature[i] = temps[kind] |
| if info is not None: |
| info.update({"fit_source": source, "fit_target": target, "fit_on": fit_on, **sizes, |
| "rows": {k: len(v) for k, v in got.items()}, |
| "report": {name: fit_report(pool, temps, before, hard) |
| for name, pool in pools.items()}}) |
| return temps |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model: DegModel, records: Path, jev_dev: Path, n_records: int = 60, |
| n_jev: int = 300) -> dict[str, Any]: |
| from deg.train import jev_example, read_jsonl |
|
|
| tok = model.renderer.tokenizer |
| report: dict[str, Any] = {} |
| for split in ("val", "test"): |
| recs = read_jsonl(records, split, n_records) |
| m: dict[str, Any] = {"records": len(recs), "valid": 0, "errors": {}} |
| dec_pairs: list[tuple[float, bool]] = [] |
| ext_ok = ext_n = cnt_ok = cnt_n = gen_fields = 0 |
| ext_ok_al = ext_n_al = 0 |
| gen_tokens = baseline_tokens = 0 |
| t0 = time.time() |
| for rec in recs: |
| plan = compile_schema(rec["schema"]) |
| baseline_tokens += len(tok(json.dumps(rec["value"], ensure_ascii=False), |
| add_special_tokens=False)["input_ids"]) |
| try: |
| resp = fill(model, rec["context"], rec["schema"], FillOptions(), plan=plan) |
| except FillError as exc: |
| key = str(exc)[:80] |
| m["errors"][key] = m["errors"].get(key, 0) + 1 |
| continue |
| m["valid"] += 1 |
| gen_tokens += resp["usage"]["generated_tokens"] |
| gold, pred = _flatten(rec["value"]), _flatten(resp["value"]) |
| modes = {} |
| for f in plan.fields: |
| modes[f.pointer] = f.mode |
| for ptr, gv in gold.items(): |
| if ptr.endswith("#len"): |
| cnt_n += 1 |
| cnt_ok += pred.get(ptr) == gv |
| continue |
| mode = modes.get(_abstract(ptr), "") |
| if mode.startswith("decide."): |
| meta = resp["meta"].get(ptr) |
| if meta is None: |
| continue |
| p = meta["probability"] |
| conf = max(p, 1 - p) if mode == "decide.noul" else p |
| dec_pairs.append((conf, pred.get(ptr) == gv and ptr in pred)) |
| elif mode.startswith("extract."): |
| ext_n += 1 |
| ext_ok += ptr in pred and pred[ptr] == gv |
| elif mode == "generate": |
| gen_fields += 1 |
| |
| for ptr, pv in pred.items(): |
| mode = modes.get(_abstract(ptr), "") |
| if mode.startswith("extract.") and ptr not in gold: |
| ext_n += 1 |
| ext_ok += pv is None |
| |
| aligned = _flatten(align_arrays(rec["value"], resp["value"])) |
| for ptr in set(gold) | set(aligned): |
| if not modes.get(_abstract(ptr), "").startswith("extract."): |
| continue |
| if ptr in gold: |
| ext_n_al += 1 |
| ext_ok_al += ptr in aligned and aligned[ptr] == gold[ptr] |
| else: |
| ext_n_al += 1 |
| ext_ok_al += aligned[ptr] is None |
| m.update({ |
| "valid_rate": round(m["valid"] / max(1, len(recs)), 4), |
| "decide_n": len(dec_pairs), |
| "decide_acc": round(sum(ok for _, ok in dec_pairs) / max(1, len(dec_pairs)), 4), |
| "decide_ece": ece(dec_pairs), |
| "extract_n": ext_n, "extract_em": round(ext_ok / max(1, ext_n), 4), |
| "extract_em_aligned": round(ext_ok_al / max(1, ext_n_al), 4), |
| "count_n": cnt_n, "count_acc": round(cnt_ok / max(1, cnt_n), 4), |
| "generate_fields": gen_fields, |
| "generated_tokens": gen_tokens, "full_json_tokens": baseline_tokens, |
| "generated_vs_full_json": round(gen_tokens / max(1, baseline_tokens), 4), |
| "s_per_record": round((time.time() - t0) / max(1, len(recs)), 3), |
| }) |
| report[split] = m |
| if jev_dev.exists() and n_jev: |
| rng = random.Random(0) |
| examples = [e for r in read_jsonl(jev_dev, None, n_jev) |
| if (e := jev_example(model, r, rng, shuffle_p=0.0)) is not None] |
| got = _collect_logits(model, examples) |
| out: dict[str, Any] = {} |
| pairs_all = [] |
| for kind in ("choice", "score", "noul"): |
| pairs = [] |
| for logits, it in got[kind]: |
| probs = torch.softmax(logits / model.heads.temp(kind), -1) |
| pairs.append((float(probs.max()), |
| int(probs.argmax()) == max(range(len(it.target)), |
| key=it.target.__getitem__))) |
| pairs_all += pairs |
| out[kind] = {"n": len(pairs), |
| "acc": round(sum(ok for _, ok in pairs) / max(1, len(pairs)), 4), |
| "ece": ece(pairs)} |
| out["all"] = {"n": len(pairs_all), |
| "acc": round(sum(ok for _, ok in pairs_all) / max(1, len(pairs_all)), 4), |
| "ece": ece(pairs_all)} |
| report["jev_dev"] = out |
| report["temperature"] = dict(zip(KINDS, model.heads.temperature.tolist())) |
| return report |
|
|
|
|
| def write_calibrated(model: DegModel, src: Path, out: Path, |
| calibration: dict[str, Any]) -> None: |
| """Copy checkpoint ``src`` to ``out`` with ``model``'s fitted temperatures. Every other |
| tensor is copied from ``src`` byte for byte (a run-1 checkpoint has no anchor heads, so |
| saving ``model`` itself would add freshly initialized ones and flip it to anchor mode).""" |
| from safetensors.torch import load_file, save_file |
|
|
| src, out = Path(src), Path(out) |
| if out.resolve() == src.resolve(): |
| raise ValueError("calibrate --out must differ from --ckpt (the source is kept as is)") |
| out.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(src / "backbone.safetensors", out / "backbone.safetensors") |
| heads = load_file(str(src / "heads.safetensors")) |
| heads["temperature"] = model.heads.temperature.detach().float().cpu().contiguous() |
| save_file(heads, str(out / "heads.safetensors")) |
| meta = json.loads((src / "deg.json").read_text()) |
| meta["temperature"] = dict(zip(KINDS, model.heads.temperature.tolist())) |
| meta["calibration"] = {"source": str(src), **calibration} |
| (out / "deg.json").write_text(json.dumps(meta, indent=2)) |
|
|
|
|
| def calibrate_main(argv: list[str]) -> None: |
| from deg.train import DEG_RECORDS, JEV_DEV |
|
|
| ap = argparse.ArgumentParser(description="fit head temperatures; write a calibrated copy") |
| ap.add_argument("--ckpt", type=Path, required=True) |
| ap.add_argument("--out", type=Path, required=True) |
| ap.add_argument("--records", type=Path, action="append", default=None, |
| help=f"datagen records file(s) whose val split is fit on " |
| f"(repeatable; default {DEG_RECORDS})") |
| ap.add_argument("--source", choices=CALIBRATION_SOURCES, default="datagen-val", |
| help="fit pool (default: datagen val only)") |
| ap.add_argument("--report-jev", action="store_true", |
| help="also report jev dev NLL/ECE before/after (separately; not fit on " |
| "unless --source includes jev-dev)") |
| ap.add_argument("--target", choices=CALIBRATION_TARGETS, default="hard", |
| help="decide NLL target: hard = one-hot filled value (default), soft = " |
| "the soft training targets") |
| ap.add_argument("--jev-dev", type=Path, default=JEV_DEV) |
| ap.add_argument("--n-jev", type=int, default=1000) |
| args = ap.parse_args(argv) |
| records = args.records or [DEG_RECORDS] |
| model = load_model(args.ckpt, device="cuda") |
| t0 = time.time() |
| info: dict[str, Any] = {} |
| temps = calibrate(model, records, args.jev_dev, args.n_jev, info=info, |
| source=args.source, report_jev=args.report_jev, target=args.target) |
| calibration = {"temperature": temps, "records": [str(r) for r in records], |
| "jev_dev": str(args.jev_dev), **info, "array_mode": model.array_mode, |
| "seconds": round(time.time() - t0, 1)} |
| write_calibrated(model, args.ckpt, args.out, calibration) |
| print(json.dumps({"out": str(args.out), **calibration})) |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| import sys |
|
|
| from deg.train import DEG_RECORDS, JEV_DEV |
|
|
| argv = sys.argv[1:] if argv is None else argv |
| if argv[:1] == ["calibrate"]: |
| calibrate_main(argv[1:]) |
| return |
|
|
| ap = argparse.ArgumentParser(description="deg eval") |
| ap.add_argument("--ckpt", type=Path, default=None, help="None = untrained heads") |
| ap.add_argument("--records", type=Path, default=DEG_RECORDS) |
| ap.add_argument("--n-records", type=int, default=60) |
| ap.add_argument("--n-jev", type=int, default=300) |
| ap.add_argument("--out", type=Path, default=None) |
| ap.add_argument("--set", choices=("datagen", "extractbench"), default="datagen", |
| help="extractbench: human-gold ExtractBench records (deg.extractbench " |
| "prepare); EVAL ONLY") |
| ap.add_argument("--extractbench-records", type=Path, default=None) |
| ap.add_argument("--predictions", type=Path, default=None, |
| help="extractbench: write one {uid, gold, pred, error} line per document") |
| args = ap.parse_args(argv) |
| model = load_model(args.ckpt, device="cuda") |
| if args.set == "extractbench": |
| from deg.extractbench import RECORDS, evaluate_records |
|
|
| report = evaluate_records(model, args.extractbench_records or RECORDS, |
| predictions=args.predictions) |
| else: |
| report = evaluate(model, args.records, JEV_DEV, args.n_records, args.n_jev) |
| text = json.dumps(report, indent=2) |
| print(text) |
| if args.out: |
| args.out.write_text(text) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
