| """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 |
| GRID = tuple(round(-3 + 0.1 * i, 2) for i in range(131)) |
| HALLUCINATION_COST = 2.0 |
| MAX_OMISSION = 0.5 |
|
|
|
|
| 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()] |
|
|
|
|
| |
|
|
|
|
| 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 |
|
|
|
|
| |
|
|
|
|
| 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:]: |
| 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: |
| 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) |
|
|
|
|
| |
|
|
|
|
| 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, |
| |
| "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 |
| 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) |
|
|
|
|
| |
|
|
|
|
| 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() |
|
|
