| """ExtractBench (Contextual AI, arXiv 2602.12247, MIT) as a deg EVAL set. Human gold: |
| evaluation only, never training data. |
| |
| ``prepare`` maps each of the 5 schemas through :func:`deg.normalize.normalize_schema` |
| (no leaf cap; ``evaluation_config`` and other non-deg keys are dropped, fields deg |
| cannot express are dropped), reports how many schemas/leaves compile, and writes one record |
| per document ``{"uid", "split": "eval", "domain", "context", "schema", "value", |
| "eval_config"}``: ``context`` is the PDF text (``<name>.txt`` next to the PDF, made with |
| PyMuPDF), ``value`` the gold projected onto the normalized schema. Image-only PDFs (no text |
| layer) are skipped and counted. |
| |
| uv run --package deg python -m deg.extractbench \\ |
| --root /data/deg/external/extractbench/dataset \\ |
| --out /data/deg/extractbench/records.jsonl |
| |
| PDF -> text first (one-time; PyMuPDF is not a deg dependency, AGPL, used as a tool): |
| |
| uv run --package deg --with pymupdf python -m deg.extractbench --pdf-text |
| |
| Scoring (:func:`score_record`) compares extract fields only (decide/generate fields and |
| array counts are reported separately): ``exact`` is value equality, ``normalized`` is |
| equality after case/whitespace/punctuation folding (numbers compare as floats). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| import unicodedata |
| from pathlib import Path |
| from typing import Any |
|
|
| from deg.normalize import count_leaves, normalize_schema |
| from deg.schema import SchemaError, compile_schema |
|
|
| ROOT = Path("/data/deg/external/extractbench/dataset") |
| RECORDS = Path("/data/deg/extractbench/records.jsonl") |
|
|
|
|
| def load_schema(path: Path) -> dict[str, Any]: |
| doc = json.loads(path.read_text()) |
| return doc["schema_definition"] if "schema_definition" in doc else doc |
|
|
|
|
| def raw_leaves(schema: Any, root: Any = None, ptr: str = "", |
| seen: tuple[str, ...] = ()) -> dict[str, str | None]: |
| """Abstract pointer -> ExtractBench ``evaluation_config`` (or None) of every leaf of the |
| raw schema (local ``$ref`` followed, unions walked at the same pointer).""" |
| root = schema if root is None else root |
| if not isinstance(schema, dict): |
| return {} |
| ref = schema.get("$ref") |
| if isinstance(ref, str) and ref.startswith("#/") and ref not in seen: |
| target: Any = root |
| for part in ref[2:].split("/"): |
| target = target.get(part, {}) if isinstance(target, dict) else {} |
| return raw_leaves({**target, **{k: v for k, v in schema.items() if k != "$ref"}}, |
| root, ptr, (*seen, ref)) |
| out: dict[str, str | None] = {} |
| branches = [b for k in ("anyOf", "oneOf", "allOf") for b in schema.get(k) or []] |
| for name, sub in (schema.get("properties") or {}).items(): |
| out |= raw_leaves(sub, root, f"{ptr}/{name}", seen) |
| if isinstance(schema.get("items"), dict): |
| out |= raw_leaves(schema["items"], root, f"{ptr}/*", seen) |
| for branch in branches: |
| out |= raw_leaves(branch, root, ptr, seen) |
| if not out and ptr: |
| config = schema.get("evaluation_config") |
| out[ptr] = config if isinstance(config, str) else None |
| return out |
|
|
|
|
| def project(value: Any, schema: dict[str, Any]) -> Any: |
| """Gold restricted to the fields the normalized schema kept.""" |
| if schema.get("type") == "object": |
| if not isinstance(value, dict): |
| return value |
| return {k: project(v, schema["properties"][k]) for k, v in value.items() |
| if k in schema["properties"]} |
| if schema.get("type") == "array" and isinstance(value, list): |
| return [project(v, schema["items"]) for v in value] |
| return value |
|
|
|
|
| def pdf_text(root: Path) -> int: |
| """Write ``<name>.txt`` (PyMuPDF text layer, pages joined by newlines) per PDF.""" |
| import pymupdf as fitz |
|
|
| n = 0 |
| for pdf in sorted(root.glob("*/*/pdf+gold/*.pdf")): |
| with fitz.open(pdf) as doc: |
| pdf.with_suffix(".txt").write_text("\n".join(page.get_text() for page in doc)) |
| n += 1 |
| return n |
|
|
|
|
| def prepare(root: Path, out: Path) -> dict[str, Any]: |
| report: dict[str, Any] = {"schemas": {}, "documents": 0, "skipped": {}} |
| rows = [] |
| for schema_path in sorted(root.glob("*/*/*-schema.json")): |
| raw = load_schema(schema_path) |
| name = schema_path.name.removesuffix("-schema.json") |
| try: |
| compile_schema(raw) |
| as_is = True |
| except SchemaError: |
| as_is = False |
| schema, why = normalize_schema(raw, max_leaves=10_000, max_depth=8) |
| info: dict[str, Any] = {"compiles_as_is": as_is, "normalized": why} |
| report["schemas"][name] = info |
| if schema is None: |
| continue |
| plan = compile_schema(schema) |
| info["leaves_kept"] = count_leaves(schema) |
| kept = {f.pointer for f in plan.fields} |
| leaves = raw_leaves(raw) |
| info["leaves_raw"] = len(leaves) |
| info["dropped"] = sorted(p for p in leaves if p not in kept) |
| info["modes"] = {m: sum(f.mode == m for f in plan.fields) |
| for m in sorted({f.mode for f in plan.fields})} |
| configs = {p: c for p, c in leaves.items() if c is not None} |
| domain = "/".join(schema_path.parent.relative_to(root).parts) |
| for pdf in sorted(schema_path.parent.glob("pdf+gold/*.pdf")): |
| gold_path = pdf.parent / (pdf.stem + ".gold.json") |
| txt = pdf.with_suffix(".txt") |
| if not gold_path.exists(): |
| report["skipped"]["no_gold"] = report["skipped"].get("no_gold", 0) + 1 |
| continue |
| if not txt.exists(): |
| report["skipped"]["no_text"] = report["skipped"].get("no_text", 0) + 1 |
| continue |
| context = txt.read_text() |
| if len(context.split()) < 20: |
| report["skipped"]["image_only"] = report["skipped"].get("image_only", 0) + 1 |
| continue |
| rows.append({"uid": f"extractbench:{domain}/{pdf.stem}", "split": "eval", |
| "domain": f"extractbench:{domain}", "context": context, |
| "schema": schema, |
| "value": project(json.loads(gold_path.read_text()), schema), |
| "eval_config": configs}) |
| out.parent.mkdir(parents=True, exist_ok=True) |
| with out.open("w", encoding="utf-8") as f: |
| for row in rows: |
| f.write(json.dumps(row, ensure_ascii=False) + "\n") |
| report["documents"] = len(rows) |
| report["words"] = sorted(len(r["context"].split()) for r in rows) |
| return report |
|
|
|
|
| _PUNCT = re.compile(r"[\W_]+", re.UNICODE) |
|
|
|
|
| def normalize_value(v: Any) -> Any: |
| if isinstance(v, bool) or v is None: |
| return v |
| if isinstance(v, (int, float)): |
| return float(v) |
| s = unicodedata.normalize("NFKC", str(v)).casefold() |
| s = _PUNCT.sub(" ", s).strip() |
| try: |
| return float(s.replace(" ", "")) if re.fullmatch(r"[\d ]+(\.\d+)?", s) else s |
| except ValueError: |
| return s |
|
|
|
|
| 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 _abstract(ptr: str) -> str: |
| return "/".join("*" if p.isdigit() else p for p in ptr.split("/")) |
|
|
|
|
| def score_record(schema: dict[str, Any], gold: Any, pred: Any, align: bool = False |
| ) -> dict[str, int]: |
| """Counts for one document: extract fields (gold non-null or predicted), exact and |
| normalized matches; array count matches. ``align``: array items are first matched to |
| their best gold item (order-insensitive, normalized equality) instead of by index.""" |
| if align: |
| from deg.evaluate import align_arrays |
|
|
| pred = align_arrays(gold, pred, lambda g, p: normalize_value(g) == normalize_value(p)) |
| modes = {f.pointer: f.mode for f in compile_schema(schema).fields} |
| g, p = _flatten(gold), _flatten(pred) |
| c = {"extract_n": 0, "exact": 0, "normalized": 0, "count_n": 0, "count_ok": 0} |
| for ptr in sorted(set(g) | set(p)): |
| if ptr.endswith("#len"): |
| if ptr in g: |
| c["count_n"] += 1 |
| c["count_ok"] += p.get(ptr) == g[ptr] |
| continue |
| if not modes.get(_abstract(ptr), "").startswith("extract."): |
| continue |
| gv, pv = g.get(ptr), p.get(ptr) |
| if gv is None and pv is None: |
| continue |
| c["extract_n"] += 1 |
| c["exact"] += gv == pv |
| c["normalized"] += normalize_value(gv) == normalize_value(pv) |
| return c |
|
|
|
|
| def evaluate_records(model: Any, records: Path, n: int | None = None, |
| predictions: Path | None = None) -> dict[str, Any]: |
| """Run :func:`deg.infer.fill` on prepared ExtractBench records; extract-field exact |
| and normalized match per domain and overall (human gold; eval only). ``predictions``: |
| also write one ``{uid, gold, pred, error}`` JSON line per document.""" |
| import gc |
| import time |
|
|
| import torch |
|
|
| from deg.infer import fill |
| from deg.request import FillOptions |
| from deg.response import FillError |
|
|
| rows = [json.loads(line) for line in records.read_text().splitlines() if line.strip()] |
| rows = rows[:n] if n else rows |
| per: dict[str, dict[str, Any]] = {} |
| docs: list[dict[str, Any]] = [] |
| pf = predictions.open("w", encoding="utf-8") if predictions else None |
| cuda = torch.cuda.is_available() and model.device.type == "cuda" |
| for row in rows: |
| m = per.setdefault(row["domain"], {"docs": 0, "valid": 0, "errors": {}, |
| "extract_n": 0, "exact": 0, "normalized": 0, |
| "count_n": 0, "count_ok": 0}) |
| m["docs"] += 1 |
| doc: dict[str, Any] = {"uid": row["uid"], "words": len(row["context"].split()), |
| "tokens": len(model.renderer.context(row["context"]).ids)} |
| docs.append(doc) |
| if cuda: |
| torch.cuda.synchronize() |
| torch.cuda.reset_peak_memory_stats() |
| base = torch.cuda.memory_allocated() |
| t0 = time.perf_counter() |
| try: |
| resp = fill(model, row["context"], row["schema"], FillOptions()) |
| err = None |
| except (FillError, torch.OutOfMemoryError) as exc: |
| |
| resp, err = None, (f"{type(exc).__name__}: {exc}" if isinstance( |
| exc, torch.OutOfMemoryError) else str(exc))[:80] |
| if cuda: |
| torch.cuda.synchronize() |
| doc["peak_gpu_mb"] = round(torch.cuda.max_memory_allocated() / 2**20) |
| doc["peak_over_weights_mb"] = round( |
| (torch.cuda.max_memory_allocated() - base) / 2**20) |
| doc["seconds"] = round(time.perf_counter() - t0, 3) |
| if pf: |
| pf.write(json.dumps({"uid": row["uid"], "gold": row["value"], |
| "pred": resp["value"] if resp else None, "error": err}, |
| ensure_ascii=False) + "\n") |
| pf.flush() |
| if resp is None: |
| doc["error"] = err |
| m["errors"][err] = m["errors"].get(err, 0) + 1 |
| if cuda: |
| gc.collect() |
| torch.cuda.empty_cache() |
| print(json.dumps(doc), flush=True) |
| continue |
| m["valid"] += 1 |
| for k, v in score_record(row["schema"], row["value"], resp["value"]).items(): |
| m[k] += v |
| doc[k] = v |
| for k, v in score_record(row["schema"], row["value"], resp["value"], align=True).items(): |
| if k in ("exact", "normalized"): |
| m[f"{k}_aligned"] = m.get(f"{k}_aligned", 0) + v |
| doc[f"{k}_aligned"] = v |
| print(json.dumps(doc), flush=True) |
| if pf: |
| pf.close() |
| total = {k: sum(m[k] for m in per.values()) |
| for k in ("docs", "valid", "extract_n", "exact", "normalized", "count_n", |
| "count_ok")} |
| for k in ("exact_aligned", "normalized_aligned"): |
| total[k] = sum(m.get(k, 0) for m in per.values()) |
| for m in [*per.values(), total]: |
| m["exact_rate"] = round(m["exact"] / max(1, m["extract_n"]), 4) |
| m["normalized_rate"] = round(m["normalized"] / max(1, m["extract_n"]), 4) |
| m["normalized_aligned_rate"] = round(m.get("normalized_aligned", 0) |
| / max(1, m["extract_n"]), 4) |
| return {"by_domain": per, "all": total, "docs": docs} |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| ap = argparse.ArgumentParser(description="ExtractBench -> deg eval records") |
| ap.add_argument("--root", type=Path, default=ROOT) |
| ap.add_argument("--out", type=Path, default=RECORDS) |
| ap.add_argument("--pdf-text", action="store_true", help="first convert PDFs to .txt") |
| args = ap.parse_args(argv) |
| if args.pdf_text: |
| print(json.dumps({"pdfs_converted": pdf_text(args.root)})) |
| print(json.dumps(prepare(args.root, args.out), indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
