Rename package to deg (DEG_ env vars, deg.json); note upcoming modalities

16d760d verified 15 days ago

22.7 kB

"""Extract null-margin calibration (EVAL ONLY): choose ``options.null_margin`` on held-out data.
A nullable/optional extract leaf is null when ``best_span_logprob < null_logprob + margin``
(:func:`deg.infer.extract_result`). The user's rule: a made-up answer costs twice an
omission, so the operating point minimises
    cost(m) = 2 * made_up_rate(unanswerable leaves) + omission_rate(answerable leaves)
on HELD-OUT data (never the deg-suite v1 test items), and must not collapse to
always-null. Steps:
    # held-out sets: SQuAD v2 dev paragraphs NOT in suite-v1 (suite contexts excluded and
    # SuiteGuard-checked) + datagen val-split records (full-off, extract-real-v1, v3-long)
    python -m deg.bench.null_margin build --out results/deg/null-margin-v1
    # one model pass per set: per top-level nullable/optional extract leaf, the best valid
    # span (value, log-prob) and the null log-prob
    CUDA_VISIBLE_DEVICES=<uuid> python -m deg.bench.null_margin dump --ckpt <ckpt> \\
        --label deg2b-1-step-0004000 --out results/deg/null-margin-v1
    # offline sweep -> curve + chosen margin per label
    python -m deg.bench.null_margin sweep --out results/deg/null-margin-v1
    # made-up / wrong-null rates of suite-v1 prediction dirs (before/after)
    python -m deg.bench.null_margin rates --system <suite label> [...]
Only top-level leaves are dumped (their pointer is unique per record); the margin itself
applies to every nullable/optional extract leaf, also inside array items. The dump probes
each leaf at margins -inf and +inf beside the real (margin 0) call, which decodes the same
head outputs, so the sweep is exact: at margin m a leaf is the best valid span when its
log-prob >= null + m, else null (checked against the margin-0 output for every leaf).
"""

from __future__ import annotations

import argparse
import dataclasses
import hashlib
import json
import math
from collections.abc import Iterable
from pathlib import Path
from typing import Any

OUT = Path("/shared/rwkv/results/deg/null-margin-v1")
SQUAD_DEV = Path("/data/deg/external/suite-v1/squad_v2/validation.parquet")
DATAGEN = Path("/data/deg/datagen")
DATAGEN_SOURCES = ("full-off", "extract-real-v1", "v3-long")
SQUAD_Q_PER_PARA = 4
HELDOUT_SQUAD_PARAS = 300
BIG = 1e30  # a probe margin: never null (-BIG) / always null (+BIG)
GRID = tuple(round(-3 + 0.1 * i, 2) for i in range(131))  # -3.0 .. 10.0 nats
HALLUCINATION_COST = 2.0
MAX_OMISSION = 0.5  # an operating point that drops half the answerable leaves is a collapse


def sha1(s: str) -> str:
    return hashlib.sha1(s.encode()).hexdigest()


def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    n = 0
    with path.open("w", encoding="utf-8") as fh:
        for r in rows:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
            n += 1
    return n


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(x) for x in path.read_text(encoding="utf-8").split("\n") if x.strip()]


# ------------------------------------------------------------------------ held-out sets


def squad_records(rows: Iterable[dict[str, Any]], exclude_contexts: set[str], guard: Any,
                  limit: int) -> tuple[list[dict[str, Any]], dict[str, int]]:
    """SQuAD v2 rows (``context, id, question, title, answers``) -> suite-format records
    (one per paragraph, up to 4 questions as nullable string fields, sha1 order), skipping
    suite paragraphs and anything the suite guard flags."""
    paras: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        paras.setdefault(row["context"], []).append(row)
    out, skipped = [], {"suite_context": 0, "guard": 0, "few_questions": 0}
    for ctx in sorted(paras, key=sha1):
        if ctx in exclude_contexts:
            skipped["suite_context"] += 1
            continue
        if guard is not None and guard.match(ctx):
            skipped["guard"] += 1
            continue
        qs = sorted(paras[ctx], key=lambda q: sha1(q["id"]))[:SQUAD_Q_PER_PARA]
        if len(qs) < 3:
            skipped["few_questions"] += 1
            continue
        props, value, alts = {}, {}, {}
        for i, q in enumerate(qs):
            key = f"q{i + 1}"
            props[key] = {"type": ["string", "null"],
                          "description": q["question"] + " (copy the answer span; null if the "
                                                         "passage does not answer it)"}
            value[key] = q["answers"][0] if q["answers"] else None
            alts[f"/{key}"] = list(q["answers"])
        out.append({"uid": "squad2-heldout:" + sha1(ctx)[:16], "context": ctx,
                    "schema": {"type": "object", "properties": props, "required": list(props),
                               "additionalProperties": False},
                    "value": value, "alternatives": alts, "domain": f"squad2/{qs[0]['title']}",
                    "question_ids": [q["id"] for q in qs]})
        if len(out) == limit:
            break
    return out, skipped


def load_squad_dev(path: Path = SQUAD_DEV) -> list[dict[str, Any]]:
    import pandas as pd

    df = pd.read_parquet(path)
    return [{"context": r.context, "id": r.id, "question": r.question, "title": r.title,
             "answers": sorted({str(a) for a in r.answers["text"]},
                               key=list(r.answers["text"]).index)}
            for r in df.itertuples()]


def nullable_top_leaves(schema: Any) -> list[str]:
    """Pointers of extract leaves outside arrays (top-level or in nested objects) that may be
    null or omitted (on_null != forbid)."""
    from deg.schema import compile_schema

    return [f.pointer for f in compile_schema(schema).fields
            if f.mode.startswith("extract.") and "*" not in f.pointer and f.on_null != "forbid"]


def datagen_records(paths: dict[str, Path], guard: Any, jev_guard: Any,
                    exclude_uids: set[str]) -> tuple[list[dict[str, Any]], dict[str, int]]:
    """``split == "val"`` records with >= 1 nullable/optional top-level extract leaf, as
    bench records (``x-generate`` fields dropped: extract outputs do not depend on them)."""
    from deg.bench.score import drop_generate

    out, skipped = [], {"guard": 0, "jev_guard": 0, "suite_uid": 0, "no_nullable_leaf": 0}
    for name, path in paths.items():
        for line in path.open(encoding="utf-8"):
            r = json.loads(line)
            if r.get("split") != "val":
                continue
            if r["uid"] in exclude_uids:
                skipped["suite_uid"] += 1
                continue
            if guard is not None and guard.match(r["context"]):
                skipped["guard"] += 1
                continue
            if jev_guard is not None and jev_guard.match(r["context"]):
                skipped["jev_guard"] += 1
                continue
            schema = drop_generate(r["schema"])
            if not nullable_top_leaves(schema):
                skipped["no_nullable_leaf"] += 1
                continue
            out.append({"uid": f"{name}:{r['uid']}", "context": r["context"], "schema": schema,
                        "value": r["value"], "domain": r.get("domain")})
    return out, skipped


def suite_uids() -> set[str]:
    from deg.bench.suite import load_source

    return {r["uid"] for r in load_source("datagen-test")}


def build(out: Path, squad_paras: int = HELDOUT_SQUAD_PARAS) -> dict[str, Any]:
    from deg.bench.jevbench import ContaminationGuard
    from deg.bench.suite import SuiteGuard, load_source

    guard, jev = SuiteGuard(), ContaminationGuard()
    suite_ctx = {r["context"] for r in load_source("squad2")}
    suite_qids = {q for r in load_source("squad2") for q in r["question_ids"]}
    squad, sq_skip = squad_records(load_squad_dev(), suite_ctx, guard, squad_paras)
    assert not {q for r in squad for q in r["question_ids"]} & suite_qids
    dg, dg_skip = datagen_records({n: DATAGEN / n / "records.jsonl" for n in DATAGEN_SOURCES},
                                  guard, jev, suite_uids())
    manifest = {"squad2-heldout": {"records": write_jsonl(out / "heldout" / "squad2.jsonl", squad),
                                   "skipped": sq_skip, "suite_guard_texts": guard.loaded,
                                   "origin": "SQuAD v2 dev paragraphs not in deg-suite v1"},
                "datagen-val": {"records": write_jsonl(out / "heldout" / "datagen-val.jsonl", dg),
                                "skipped": dg_skip,
                                "origin": "split=val records of " + ", ".join(DATAGEN_SOURCES)}}
    for name in ("squad2", "datagen-val"):
        p = out / "heldout" / f"{name}.jsonl"
        manifest["squad2-heldout" if name == "squad2" else name]["sha256"] = hashlib.sha256(
            p.read_bytes()).hexdigest()
    (out / "heldout" / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n")
    return manifest


# ------------------------------------------------------------------------------- dump


class LeafProbe:
    """Wraps :func:`deg.infer.extract_result`: beside the real call, records for each
    top-level nullable/optional leaf the best valid span (margin -inf) and the null
    outcome (margin +inf)."""

    def __init__(self) -> None:
        self.rows: list[dict[str, Any]] = []
        self._orig: Any = None

    def __enter__(self) -> LeafProbe:
        from deg import infer

        self._orig = orig = infer.extract_result

        def probed(fld: Any, ctx_text: str, ctx: Any, start: Any, end: Any, temperature: Any,
                   options: Any) -> Any:
            res = orig(fld, ctx_text, ctx, start, end, temperature, options)
            if fld.on_null != "forbid" and "*" not in fld.pointer:
                best = orig(fld, ctx_text, ctx, start, end, temperature,
                            dataclasses.replace(options, null_margin=-BIG))
                null = orig(fld, ctx_text, ctx, start, end, temperature,
                            dataclasses.replace(options, null_margin=BIG))
                span_lp = (_log(best.meta["confidence"]) if best.value is not None
                           else -math.inf)
                self.rows.append({"ptr": fld.pointer, "value": res.value,
                                  "span_value": best.value, "span_lp": span_lp,
                                  "null_lp": _log(null.meta["confidence"])})
            return res

        infer.extract_result = probed
        return self

    def __exit__(self, *exc: Any) -> None:
        from deg import infer

        infer.extract_result = self._orig


def _log(p: float) -> float:
    return math.log(p) if p > 0 else -math.inf


def leaf_gold(rec: dict[str, Any], ptr: str) -> tuple[str, list[Any]]:
    """("unans", []) when the gold is null/absent, else ("ans", accepted answers)."""
    if "alternatives" in rec and ptr in rec["alternatives"]:
        alts = rec["alternatives"][ptr]
        return ("ans", alts) if alts else ("unans", [])
    gold: Any = rec["value"]
    for part in ptr.split("/")[1:]:  # nested object leaves (no arrays: top-level dumps)
        key = part.replace("~1", "/").replace("~0", "~")
        gold = gold.get(key) if isinstance(gold, dict) else None
    return ("unans", []) if gold is None else ("ans", [gold])


def dump(ckpt: str, label: str, out: Path, sets: list[str]) -> None:
    import torch

    from deg.infer import fill
    from deg.model import load_model
    from deg.request import FillOptions

    model = load_model(Path(ckpt), device="cuda")
    for name in sets:
        recs = read_jsonl(out / "heldout" / f"{name}.jsonl")
        path = out / "dumps" / label / f"{name}.jsonl"
        path.parent.mkdir(parents=True, exist_ok=True)
        done = {r["uid"] for r in read_jsonl(path)} if path.exists() else set()
        with path.open("a", encoding="utf-8") as fh, LeafProbe() as probe:
            for i, rec in enumerate(recs):
                if rec["uid"] in done:
                    continue
                probe.rows = []
                err = None
                try:
                    fill(model, rec["context"], rec["schema"], FillOptions())
                except Exception as exc:  # noqa: BLE001 - recorded, excluded from the sweep
                    err = f"{type(exc).__name__}: {str(exc)[:160]}"
                    torch.cuda.empty_cache()
                leaves = []
                for row in probe.rows:
                    kind, alts = leaf_gold(rec, row["ptr"])
                    leaves.append({**row, "gold": kind, "alts": alts})
                fh.write(json.dumps({"uid": rec["uid"], "error": err, "leaves": leaves},
                                    ensure_ascii=False) + "\n")
                fh.flush()
                if (i + 1) % 50 == 0:
                    print(json.dumps({"label": label, "set": name, "done": i + 1,
                                      "of": len(recs)}), flush=True)
    print(json.dumps({"label": label, "dumped": sets, "ckpt": ckpt}), flush=True)


# ------------------------------------------------------------------------------ sweep


def is_null(leaf: dict[str, Any], m: float) -> bool:
    return leaf["span_value"] is None or leaf["span_lp"] < leaf["null_lp"] + m


def matches(value: Any, alts: list[Any]) -> bool:
    from deg.bench.score import normalize

    return value is not None and normalize(value) in {normalize(a) for a in alts}


def point(leaves: list[dict[str, Any]], m: float) -> dict[str, Any]:
    ans = [x for x in leaves if x["gold"] == "ans"]
    un = [x for x in leaves if x["gold"] == "unans"]
    omit = sum(is_null(x, m) for x in ans)
    right = sum(not is_null(x, m) and matches(x["span_value"], x["alts"]) for x in ans)
    made_up = sum(not is_null(x, m) for x in un)
    mu = made_up / len(un) if un else 0.0
    om = omit / len(ans) if ans else 0.0
    return {"margin": m, "answerable_n": len(ans), "unanswerable_n": len(un),
            "made_up_rate": mu, "omission_rate": om,
            "answerable_correct": right / len(ans) if ans else None,
            "answerable_wrong_span": (len(ans) - omit - right) / len(ans) if ans else None,
            "unanswerable_null": 1 - mu,
            # suite-style balanced score (squad2 component definition)
            "balanced": ((right / len(ans) if ans else 0.0) + (1 - mu)) / 2,
            "cost": HALLUCINATION_COST * mu + om}


def check_margin0(leaves: list[dict[str, Any]]) -> int:
    """Leaves whose offline margin-0 decision differs from the real margin-0 output (0 =
    the sweep reproduces inference exactly)."""
    return sum((x["value"] is None) != is_null(x, 0.0)
               or (x["value"] is not None and x["value"] != x["span_value"]) for x in leaves)


def choose(curve: list[dict[str, Any]]) -> dict[str, Any]:
    """Min cost (the mean of the per-set costs) among non-collapsed points (omission <=
    MAX_OMISSION on every set); ties -> the margin closest to 0."""
    ok = [p for p in curve if all(s["omission_rate"] <= MAX_OMISSION for s in p["sets"].values())]
    return min(ok, key=lambda p: (round(p["cost"], 9), abs(p["margin"])))


def regold(rows: list[dict[str, Any]], heldout: Path) -> list[dict[str, Any]]:
    """The leaves of error-free dump rows, gold re-derived from the held-out records (the
    dump's own gold fields are not trusted: the sweep always uses :func:`leaf_gold`)."""
    recs = {r["uid"]: r for r in read_jsonl(heldout)}
    out = []
    for r in rows:
        if r["error"] is not None:
            continue
        for leaf in r["leaves"]:
            kind, alts = leaf_gold(recs[r["uid"]], leaf["ptr"])
            out.append({**leaf, "gold": kind, "alts": alts})
    return out


def sweep(out: Path, labels: list[str] | None = None) -> dict[str, Any]:
    result: dict[str, Any] = {"rule": f"cost = {HALLUCINATION_COST:g}*made_up_rate(unanswerable)"
                                      " + omission_rate(answerable), averaged over the held-out"
                                      " sets; omission <= "
                                      f"{MAX_OMISSION:g} on every set (no always-null collapse)",
                              "models": {}}
    for d in sorted((out / "dumps").iterdir()):
        if labels and d.name not in labels:
            continue
        sets = {p.stem: regold(read_jsonl(p), out / "heldout" / p.name)
                for p in sorted(d.glob("*.jsonl"))}
        curve = []
        for m in GRID:
            per = {k: point(v, m) for k, v in sets.items()}
            curve.append({"margin": m, "cost": sum(s["cost"] for s in per.values()) / len(per),
                          "sets": per})
        best = choose(curve)
        zero = next(p for p in curve if p["margin"] == 0)
        result["models"][d.name] = {
            "chosen_margin": best["margin"], "chosen": best, "margin0": zero,
            "margin0_mismatch": {k: check_margin0(v) for k, v in sets.items()},
            "errors": {p.stem: sum(r["error"] is not None for r in read_jsonl(p))
                       for p in sorted(d.glob("*.jsonl"))},
            "curve": curve}
    (out / "sweep.json").write_text(json.dumps(result, indent=1) + "\n")
    (out / "curve.md").write_text(curve_md(result))
    return result


def curve_md(result: dict[str, Any]) -> str:
    lines = ["# null-margin curve (held-out)", "", result["rule"], ""]
    for label, m in result["models"].items():
        names = list(m["chosen"]["sets"])
        lines += [f"## {label}", "", f"chosen margin **{m['chosen_margin']:g}** nats "
                  f"(margin-0 reproduction mismatches: {m['margin0_mismatch']}; "
                  f"errors: {m['errors']})", "",
                  "| margin | cost | " + " | ".join(
                      f"{n} made-up | {n} omit | {n} ans-correct | {n} balanced"
                      for n in names) + " |",
                  "|---|---|" + "---|" * (4 * len(names))]
        for p in m["curve"]:
            if round(p["margin"] * 10) % 5 and p["margin"] != m["chosen_margin"]:
                continue  # every 0.5 nats + the chosen point
            mark = " **<-**" if p["margin"] == m["chosen_margin"] else ""
            lines.append(f"| {p['margin']:g}{mark} | {p['cost']:.3f} | " + " | ".join(
                f"{s['made_up_rate']:.3f} | {s['omission_rate']:.3f} | "
                f"{s['answerable_correct']:.3f} | {s['balanced']:.3f}"
                for s in p["sets"].values()) + " |")
        lines.append("")
    return "\n".join(lines)


# ------------------------------------------------------------------------ suite rates


def extract_null_rates(recs: list[dict[str, Any]], preds: dict[str, Any]) -> dict[str, Any]:
    """Made-up (gold null/absent, prediction non-null) and wrong-null (gold non-null,
    prediction null/absent) rates over the extract leaves of suite records. A record with an
    error counts every leaf as null (its value is missing)."""
    from deg.evaluate import _abstract, _flatten
    from deg.schema import compile_schema

    ans = un = made_up = wrong_null = 0
    for rec in recs:
        p = preds.get(rec["uid"])
        got = _flatten(p.value) if p is not None and p.value is not None else {}
        fields = {f.pointer: f for f in compile_schema(rec["schema"]).fields}
        gold = _flatten(rec["value"])
        if "alternatives" in rec:
            gold = {ptr: (alts[0] if alts else None) for ptr, alts in rec["alternatives"].items()}
        ptrs = set(gold) | {f.pointer for f in fields.values()
                            if "*" not in f.pointer and f.mode.startswith("extract.")}
        for ptr in ptrs:
            if ptr.endswith("#len"):
                continue
            f = fields.get(_abstract(ptr))
            if f is None or not f.mode.startswith("extract.") or f.on_null == "forbid":
                continue
            gv, pv = gold.get(ptr), got.get(ptr)
            if gv is None:
                un += 1
                made_up += pv is not None
            else:
                ans += 1
                wrong_null += pv is None
    return {"unanswerable_n": un, "answerable_n": ans,
            "made_up_rate": made_up / un if un else None,
            "wrong_null_rate": wrong_null / ans if ans else None}


def suite_rates(system: str, results: Path | None = None,
                sources: tuple[str, ...] = ("datagen-test", "squad2", "mixed")) -> dict[str, Any]:
    from deg.bench.run import load_preds
    from deg.bench.suite import RESULTS, load_source, system_dir

    sdir = system_dir(results or RESULTS, system)
    out = {}
    for src in sources:
        path = sdir / f"{src}.preds.jsonl"
        if path.exists():
            out[src] = extract_null_rates(load_source(src), load_preds(path, keep_errors=True))
    return out


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    b = sub.add_parser("build")
    b.add_argument("--out", type=Path, default=OUT)
    b.add_argument("--squad-paras", type=int, default=HELDOUT_SQUAD_PARAS)
    d = sub.add_parser("dump")
    d.add_argument("--ckpt", required=True)
    d.add_argument("--label", required=True)
    d.add_argument("--out", type=Path, default=OUT)
    d.add_argument("--sets", nargs="*", default=["squad2", "datagen-val"])
    d.add_argument("--max-gpu-gb", type=float, default=None)
    s = sub.add_parser("sweep")
    s.add_argument("--out", type=Path, default=OUT)
    s.add_argument("--labels", nargs="*", default=None)
    r = sub.add_parser("rates")
    r.add_argument("--system", nargs="+", required=True)
    r.add_argument("--results", type=Path, default=None)
    args = ap.parse_args(argv)
    if args.cmd == "build":
        print(json.dumps(build(args.out, args.squad_paras), indent=1))
    elif args.cmd == "dump":
        if args.max_gpu_gb:
            import torch
            total = torch.cuda.get_device_properties(0).total_memory
            torch.cuda.set_per_process_memory_fraction(min(1.0, args.max_gpu_gb * 2**30 / total))
        dump(args.ckpt, args.label, args.out, args.sets)
    elif args.cmd == "sweep":
        res = sweep(args.out, args.labels)
        print(json.dumps({k: {"chosen_margin": v["chosen_margin"],
                              "cost": v["chosen"]["cost"], "cost0": v["margin0"]["cost"],
                              "mismatch": v["margin0_mismatch"]}
                          for k, v in res["models"].items()}, indent=1))
    else:
        print(json.dumps({s: suite_rates(s, args.results) for s in args.system}, indent=1))


if __name__ == "__main__":
    main()

	"""Extract null-margin calibration (EVAL ONLY): choose ``options.null_margin`` on held-out data.

	A nullable/optional extract leaf is null when ``best_span_logprob < null_logprob + margin``
	(:func:`deg.infer.extract_result`). The user's rule: a made-up answer costs twice an
	omission, so the operating point minimises

	cost(m) = 2 * made_up_rate(unanswerable leaves) + omission_rate(answerable leaves)

	on HELD-OUT data (never the deg-suite v1 test items), and must not collapse to
	always-null. Steps:

	# held-out sets: SQuAD v2 dev paragraphs NOT in suite-v1 (suite contexts excluded and
	# SuiteGuard-checked) + datagen val-split records (full-off, extract-real-v1, v3-long)
	python -m deg.bench.null_margin build --out results/deg/null-margin-v1
	# one model pass per set: per top-level nullable/optional extract leaf, the best valid
	# span (value, log-prob) and the null log-prob
	CUDA_VISIBLE_DEVICES=<uuid> python -m deg.bench.null_margin dump --ckpt <ckpt> \\
	--label deg2b-1-step-0004000 --out results/deg/null-margin-v1
	# offline sweep -> curve + chosen margin per label
	python -m deg.bench.null_margin sweep --out results/deg/null-margin-v1
	# made-up / wrong-null rates of suite-v1 prediction dirs (before/after)
	python -m deg.bench.null_margin rates --system <suite label> [...]

	Only top-level leaves are dumped (their pointer is unique per record); the margin itself
	applies to every nullable/optional extract leaf, also inside array items. The dump probes
	each leaf at margins -inf and +inf beside the real (margin 0) call, which decodes the same
	head outputs, so the sweep is exact: at margin m a leaf is the best valid span when its
	log-prob >= null + m, else null (checked against the margin-0 output for every leaf).
	"""

	from __future__ import annotations

	import argparse
	import dataclasses
	import hashlib
	import json
	import math
	from collections.abc import Iterable
	from pathlib import Path
	from typing import Any

	OUT = Path("/shared/rwkv/results/deg/null-margin-v1")
	SQUAD_DEV = Path("/data/deg/external/suite-v1/squad_v2/validation.parquet")
	DATAGEN = Path("/data/deg/datagen")
	DATAGEN_SOURCES = ("full-off", "extract-real-v1", "v3-long")
	SQUAD_Q_PER_PARA = 4
	HELDOUT_SQUAD_PARAS = 300
	BIG = 1e30 # a probe margin: never null (-BIG) / always null (+BIG)
	GRID = tuple(round(-3 + 0.1 * i, 2) for i in range(131)) # -3.0 .. 10.0 nats
	HALLUCINATION_COST = 2.0
	MAX_OMISSION = 0.5 # an operating point that drops half the answerable leaves is a collapse


	def sha1(s: str) -> str:
	return hashlib.sha1(s.encode()).hexdigest()


	def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int:
	path.parent.mkdir(parents=True, exist_ok=True)
	n = 0
	with path.open("w", encoding="utf-8") as fh:
	for r in rows:
	fh.write(json.dumps(r, ensure_ascii=False) + "\n")
	n += 1
	return n


	def read_jsonl(path: Path) -> list[dict[str, Any]]:
	return [json.loads(x) for x in path.read_text(encoding="utf-8").split("\n") if x.strip()]


	# ------------------------------------------------------------------------ held-out sets


	def squad_records(rows: Iterable[dict[str, Any]], exclude_contexts: set[str], guard: Any,
	limit: int) -> tuple[list[dict[str, Any]], dict[str, int]]:
	"""SQuAD v2 rows (``context, id, question, title, answers``) -> suite-format records
	(one per paragraph, up to 4 questions as nullable string fields, sha1 order), skipping
	suite paragraphs and anything the suite guard flags."""
	paras: dict[str, list[dict[str, Any]]] = {}
	for row in rows:
	paras.setdefault(row["context"], []).append(row)
	out, skipped = [], {"suite_context": 0, "guard": 0, "few_questions": 0}
	for ctx in sorted(paras, key=sha1):
	if ctx in exclude_contexts:
	skipped["suite_context"] += 1
	continue
	if guard is not None and guard.match(ctx):
	skipped["guard"] += 1
	continue
	qs = sorted(paras[ctx], key=lambda q: sha1(q["id"]))[:SQUAD_Q_PER_PARA]
	if len(qs) < 3:
	skipped["few_questions"] += 1
	continue
	props, value, alts = {}, {}, {}
	for i, q in enumerate(qs):
	key = f"q{i + 1}"
	props[key] = {"type": ["string", "null"],
	"description": q["question"] + " (copy the answer span; null if the "
	"passage does not answer it)"}
	value[key] = q["answers"][0] if q["answers"] else None
	alts[f"/{key}"] = list(q["answers"])
	out.append({"uid": "squad2-heldout:" + sha1(ctx)[:16], "context": ctx,
	"schema": {"type": "object", "properties": props, "required": list(props),
	"additionalProperties": False},
	"value": value, "alternatives": alts, "domain": f"squad2/{qs[0]['title']}",
	"question_ids": [q["id"] for q in qs]})
	if len(out) == limit:
	break
	return out, skipped


	def load_squad_dev(path: Path = SQUAD_DEV) -> list[dict[str, Any]]:
	import pandas as pd

	df = pd.read_parquet(path)
	return [{"context": r.context, "id": r.id, "question": r.question, "title": r.title,
	"answers": sorted({str(a) for a in r.answers["text"]},
	key=list(r.answers["text"]).index)}
	for r in df.itertuples()]


	def nullable_top_leaves(schema: Any) -> list[str]:
	"""Pointers of extract leaves outside arrays (top-level or in nested objects) that may be
	null or omitted (on_null != forbid)."""
	from deg.schema import compile_schema

	return [f.pointer for f in compile_schema(schema).fields
	if f.mode.startswith("extract.") and "*" not in f.pointer and f.on_null != "forbid"]


	def datagen_records(paths: dict[str, Path], guard: Any, jev_guard: Any,
	exclude_uids: set[str]) -> tuple[list[dict[str, Any]], dict[str, int]]:
	"""``split == "val"`` records with >= 1 nullable/optional top-level extract leaf, as
	bench records (``x-generate`` fields dropped: extract outputs do not depend on them)."""
	from deg.bench.score import drop_generate

	out, skipped = [], {"guard": 0, "jev_guard": 0, "suite_uid": 0, "no_nullable_leaf": 0}
	for name, path in paths.items():
	for line in path.open(encoding="utf-8"):
	r = json.loads(line)
	if r.get("split") != "val":
	continue
	if r["uid"] in exclude_uids:
	skipped["suite_uid"] += 1
	continue
	if guard is not None and guard.match(r["context"]):
	skipped["guard"] += 1
	continue
	if jev_guard is not None and jev_guard.match(r["context"]):
	skipped["jev_guard"] += 1
	continue
	schema = drop_generate(r["schema"])
	if not nullable_top_leaves(schema):
	skipped["no_nullable_leaf"] += 1
	continue
	out.append({"uid": f"{name}:{r['uid']}", "context": r["context"], "schema": schema,
	"value": r["value"], "domain": r.get("domain")})
	return out, skipped


	def suite_uids() -> set[str]:
	from deg.bench.suite import load_source

	return {r["uid"] for r in load_source("datagen-test")}


	def build(out: Path, squad_paras: int = HELDOUT_SQUAD_PARAS) -> dict[str, Any]:
	from deg.bench.jevbench import ContaminationGuard
	from deg.bench.suite import SuiteGuard, load_source

	guard, jev = SuiteGuard(), ContaminationGuard()
	suite_ctx = {r["context"] for r in load_source("squad2")}
	suite_qids = {q for r in load_source("squad2") for q in r["question_ids"]}
	squad, sq_skip = squad_records(load_squad_dev(), suite_ctx, guard, squad_paras)
	assert not {q for r in squad for q in r["question_ids"]} & suite_qids
	dg, dg_skip = datagen_records({n: DATAGEN / n / "records.jsonl" for n in DATAGEN_SOURCES},
	guard, jev, suite_uids())
	manifest = {"squad2-heldout": {"records": write_jsonl(out / "heldout" / "squad2.jsonl", squad),
	"skipped": sq_skip, "suite_guard_texts": guard.loaded,
	"origin": "SQuAD v2 dev paragraphs not in deg-suite v1"},
	"datagen-val": {"records": write_jsonl(out / "heldout" / "datagen-val.jsonl", dg),
	"skipped": dg_skip,
	"origin": "split=val records of " + ", ".join(DATAGEN_SOURCES)}}
	for name in ("squad2", "datagen-val"):
	p = out / "heldout" / f"{name}.jsonl"
	manifest["squad2-heldout" if name == "squad2" else name]["sha256"] = hashlib.sha256(
	p.read_bytes()).hexdigest()
	(out / "heldout" / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n")
	return manifest


	# ------------------------------------------------------------------------------- dump


	class LeafProbe:
	"""Wraps :func:`deg.infer.extract_result`: beside the real call, records for each
	top-level nullable/optional leaf the best valid span (margin -inf) and the null
	outcome (margin +inf)."""

	def __init__(self) -> None:
	self.rows: list[dict[str, Any]] = []
	self._orig: Any = None

	def __enter__(self) -> LeafProbe:
	from deg import infer

	self._orig = orig = infer.extract_result

	def probed(fld: Any, ctx_text: str, ctx: Any, start: Any, end: Any, temperature: Any,
	options: Any) -> Any:
	res = orig(fld, ctx_text, ctx, start, end, temperature, options)
	if fld.on_null != "forbid" and "*" not in fld.pointer:
	best = orig(fld, ctx_text, ctx, start, end, temperature,
	dataclasses.replace(options, null_margin=-BIG))
	null = orig(fld, ctx_text, ctx, start, end, temperature,
	dataclasses.replace(options, null_margin=BIG))
	span_lp = (_log(best.meta["confidence"]) if best.value is not None
	else -math.inf)
	self.rows.append({"ptr": fld.pointer, "value": res.value,
	"span_value": best.value, "span_lp": span_lp,
	"null_lp": _log(null.meta["confidence"])})
	return res

	infer.extract_result = probed
	return self

	def __exit__(self, *exc: Any) -> None:
	from deg import infer

	infer.extract_result = self._orig


	def _log(p: float) -> float:
	return math.log(p) if p > 0 else -math.inf


	def leaf_gold(rec: dict[str, Any], ptr: str) -> tuple[str, list[Any]]:
	"""("unans", []) when the gold is null/absent, else ("ans", accepted answers)."""
	if "alternatives" in rec and ptr in rec["alternatives"]:
	alts = rec["alternatives"][ptr]
	return ("ans", alts) if alts else ("unans", [])
	gold: Any = rec["value"]
	for part in ptr.split("/")[1:]: # nested object leaves (no arrays: top-level dumps)
	key = part.replace("~1", "/").replace("~0", "~")
	gold = gold.get(key) if isinstance(gold, dict) else None
	return ("unans", []) if gold is None else ("ans", [gold])


	def dump(ckpt: str, label: str, out: Path, sets: list[str]) -> None:
	import torch

	from deg.infer import fill
	from deg.model import load_model
	from deg.request import FillOptions

	model = load_model(Path(ckpt), device="cuda")
	for name in sets:
	recs = read_jsonl(out / "heldout" / f"{name}.jsonl")
	path = out / "dumps" / label / f"{name}.jsonl"
	path.parent.mkdir(parents=True, exist_ok=True)
	done = {r["uid"] for r in read_jsonl(path)} if path.exists() else set()
	with path.open("a", encoding="utf-8") as fh, LeafProbe() as probe:
	for i, rec in enumerate(recs):
	if rec["uid"] in done:
	continue
	probe.rows = []
	err = None
	try:
	fill(model, rec["context"], rec["schema"], FillOptions())
	except Exception as exc: # noqa: BLE001 - recorded, excluded from the sweep
	err = f"{type(exc).__name__}: {str(exc)[:160]}"
	torch.cuda.empty_cache()
	leaves = []
	for row in probe.rows:
	kind, alts = leaf_gold(rec, row["ptr"])
	leaves.append({**row, "gold": kind, "alts": alts})
	fh.write(json.dumps({"uid": rec["uid"], "error": err, "leaves": leaves},
	ensure_ascii=False) + "\n")
	fh.flush()
	if (i + 1) % 50 == 0:
	print(json.dumps({"label": label, "set": name, "done": i + 1,
	"of": len(recs)}), flush=True)
	print(json.dumps({"label": label, "dumped": sets, "ckpt": ckpt}), flush=True)


	# ------------------------------------------------------------------------------ sweep


	def is_null(leaf: dict[str, Any], m: float) -> bool:
	return leaf["span_value"] is None or leaf["span_lp"] < leaf["null_lp"] + m


	def matches(value: Any, alts: list[Any]) -> bool:
	from deg.bench.score import normalize

	return value is not None and normalize(value) in {normalize(a) for a in alts}


	def point(leaves: list[dict[str, Any]], m: float) -> dict[str, Any]:
	ans = [x for x in leaves if x["gold"] == "ans"]
	un = [x for x in leaves if x["gold"] == "unans"]
	omit = sum(is_null(x, m) for x in ans)
	right = sum(not is_null(x, m) and matches(x["span_value"], x["alts"]) for x in ans)
	made_up = sum(not is_null(x, m) for x in un)
	mu = made_up / len(un) if un else 0.0
	om = omit / len(ans) if ans else 0.0
	return {"margin": m, "answerable_n": len(ans), "unanswerable_n": len(un),
	"made_up_rate": mu, "omission_rate": om,
	"answerable_correct": right / len(ans) if ans else None,
	"answerable_wrong_span": (len(ans) - omit - right) / len(ans) if ans else None,
	"unanswerable_null": 1 - mu,
	# suite-style balanced score (squad2 component definition)
	"balanced": ((right / len(ans) if ans else 0.0) + (1 - mu)) / 2,
	"cost": HALLUCINATION_COST * mu + om}


	def check_margin0(leaves: list[dict[str, Any]]) -> int:
	"""Leaves whose offline margin-0 decision differs from the real margin-0 output (0 =
	the sweep reproduces inference exactly)."""
	return sum((x["value"] is None) != is_null(x, 0.0)
	or (x["value"] is not None and x["value"] != x["span_value"]) for x in leaves)


	def choose(curve: list[dict[str, Any]]) -> dict[str, Any]:
	"""Min cost (the mean of the per-set costs) among non-collapsed points (omission <=
	MAX_OMISSION on every set); ties -> the margin closest to 0."""
	ok = [p for p in curve if all(s["omission_rate"] <= MAX_OMISSION for s in p["sets"].values())]
	return min(ok, key=lambda p: (round(p["cost"], 9), abs(p["margin"])))


	def regold(rows: list[dict[str, Any]], heldout: Path) -> list[dict[str, Any]]:
	"""The leaves of error-free dump rows, gold re-derived from the held-out records (the
	dump's own gold fields are not trusted: the sweep always uses :func:`leaf_gold`)."""
	recs = {r["uid"]: r for r in read_jsonl(heldout)}
	out = []
	for r in rows:
	if r["error"] is not None:
	continue
	for leaf in r["leaves"]:
	kind, alts = leaf_gold(recs[r["uid"]], leaf["ptr"])
	out.append({**leaf, "gold": kind, "alts": alts})
	return out


	def sweep(out: Path, labels: list[str] \| None = None) -> dict[str, Any]:
	result: dict[str, Any] = {"rule": f"cost = {HALLUCINATION_COST:g}*made_up_rate(unanswerable)"
	" + omission_rate(answerable), averaged over the held-out"
	" sets; omission <= "
	f"{MAX_OMISSION:g} on every set (no always-null collapse)",
	"models": {}}
	for d in sorted((out / "dumps").iterdir()):
	if labels and d.name not in labels:
	continue
	sets = {p.stem: regold(read_jsonl(p), out / "heldout" / p.name)
	for p in sorted(d.glob("*.jsonl"))}
	curve = []
	for m in GRID:
	per = {k: point(v, m) for k, v in sets.items()}
	curve.append({"margin": m, "cost": sum(s["cost"] for s in per.values()) / len(per),
	"sets": per})
	best = choose(curve)
	zero = next(p for p in curve if p["margin"] == 0)
	result["models"][d.name] = {
	"chosen_margin": best["margin"], "chosen": best, "margin0": zero,
	"margin0_mismatch": {k: check_margin0(v) for k, v in sets.items()},
	"errors": {p.stem: sum(r["error"] is not None for r in read_jsonl(p))
	for p in sorted(d.glob("*.jsonl"))},
	"curve": curve}
	(out / "sweep.json").write_text(json.dumps(result, indent=1) + "\n")
	(out / "curve.md").write_text(curve_md(result))
	return result


	def curve_md(result: dict[str, Any]) -> str:
	lines = ["# null-margin curve (held-out)", "", result["rule"], ""]
	for label, m in result["models"].items():
	names = list(m["chosen"]["sets"])
	lines += [f"## {label}", "", f"chosen margin {m['chosen_margin']:g} nats "
	f"(margin-0 reproduction mismatches: {m['margin0_mismatch']}; "
	f"errors: {m['errors']})", "",
	"\| margin \| cost \| " + " \| ".join(
	f"{n} made-up \| {n} omit \| {n} ans-correct \| {n} balanced"
	for n in names) + " \|",
	"\|---\|---\|" + "---\|" * (4 * len(names))]
	for p in m["curve"]:
	if round(p["margin"] * 10) % 5 and p["margin"] != m["chosen_margin"]:
	continue # every 0.5 nats + the chosen point
	mark = " <-" if p["margin"] == m["chosen_margin"] else ""
	lines.append(f"\| {p['margin']:g}{mark} \| {p['cost']:.3f} \| " + " \| ".join(
	f"{s['made_up_rate']:.3f} \| {s['omission_rate']:.3f} \| "
	f"{s['answerable_correct']:.3f} \| {s['balanced']:.3f}"
	for s in p["sets"].values()) + " \|")
	lines.append("")
	return "\n".join(lines)


	# ------------------------------------------------------------------------ suite rates


	def extract_null_rates(recs: list[dict[str, Any]], preds: dict[str, Any]) -> dict[str, Any]:
	"""Made-up (gold null/absent, prediction non-null) and wrong-null (gold non-null,
	prediction null/absent) rates over the extract leaves of suite records. A record with an
	error counts every leaf as null (its value is missing)."""
	from deg.evaluate import _abstract, _flatten
	from deg.schema import compile_schema

	ans = un = made_up = wrong_null = 0
	for rec in recs:
	p = preds.get(rec["uid"])
	got = _flatten(p.value) if p is not None and p.value is not None else {}
	fields = {f.pointer: f for f in compile_schema(rec["schema"]).fields}
	gold = _flatten(rec["value"])
	if "alternatives" in rec:
	gold = {ptr: (alts[0] if alts else None) for ptr, alts in rec["alternatives"].items()}
	ptrs = set(gold) \| {f.pointer for f in fields.values()
	if "*" not in f.pointer and f.mode.startswith("extract.")}
	for ptr in ptrs:
	if ptr.endswith("#len"):
	continue
	f = fields.get(_abstract(ptr))
	if f is None or not f.mode.startswith("extract.") or f.on_null == "forbid":
	continue
	gv, pv = gold.get(ptr), got.get(ptr)
	if gv is None:
	un += 1
	made_up += pv is not None
	else:
	ans += 1
	wrong_null += pv is None
	return {"unanswerable_n": un, "answerable_n": ans,
	"made_up_rate": made_up / un if un else None,
	"wrong_null_rate": wrong_null / ans if ans else None}


	def suite_rates(system: str, results: Path \| None = None,
	sources: tuple[str, ...] = ("datagen-test", "squad2", "mixed")) -> dict[str, Any]:
	from deg.bench.run import load_preds
	from deg.bench.suite import RESULTS, load_source, system_dir

	sdir = system_dir(results or RESULTS, system)
	out = {}
	for src in sources:
	path = sdir / f"{src}.preds.jsonl"
	if path.exists():
	out[src] = extract_null_rates(load_source(src), load_preds(path, keep_errors=True))
	return out


	def main(argv: list[str] \| None = None) -> None:
	ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
	sub = ap.add_subparsers(dest="cmd", required=True)
	b = sub.add_parser("build")
	b.add_argument("--out", type=Path, default=OUT)
	b.add_argument("--squad-paras", type=int, default=HELDOUT_SQUAD_PARAS)
	d = sub.add_parser("dump")
	d.add_argument("--ckpt", required=True)
	d.add_argument("--label", required=True)
	d.add_argument("--out", type=Path, default=OUT)
	d.add_argument("--sets", nargs="*", default=["squad2", "datagen-val"])
	d.add_argument("--max-gpu-gb", type=float, default=None)
	s = sub.add_parser("sweep")
	s.add_argument("--out", type=Path, default=OUT)
	s.add_argument("--labels", nargs="*", default=None)
	r = sub.add_parser("rates")
	r.add_argument("--system", nargs="+", required=True)
	r.add_argument("--results", type=Path, default=None)
	args = ap.parse_args(argv)
	if args.cmd == "build":
	print(json.dumps(build(args.out, args.squad_paras), indent=1))
	elif args.cmd == "dump":
	if args.max_gpu_gb:
	import torch
	total = torch.cuda.get_device_properties(0).total_memory
	torch.cuda.set_per_process_memory_fraction(min(1.0, args.max_gpu_gb * 2**30 / total))
	dump(args.ckpt, args.label, args.out, args.sets)
	elif args.cmd == "sweep":
	res = sweep(args.out, args.labels)
	print(json.dumps({k: {"chosen_margin": v["chosen_margin"],
	"cost": v["chosen"]["cost"], "cost0": v["margin0"]["cost"],
	"mismatch": v["margin0_mismatch"]}
	for k, v in res["models"].items()}, indent=1))
	else:
	print(json.dumps({s: suite_rates(s, args.results) for s in args.system}, indent=1))


	if __name__ == "__main__":
	main()