spb/wpai0002_uqo

1#!/usr/bin/env python32"""Generate reconstructions for the pairs of a split with a checkpoint (and/or baselines) → outputs JSONL."""3from __future__ import annotations45import argparse6import sys7import time8from pathlib import Path910ROOT = Path(__file__).resolve().parents[1]11sys.path.insert(0, str(ROOT / "src"))12from htr.evaluation.baselines import BASELINES  # noqa: E40213from htr.utils.io import read_jsonl, write_json, write_jsonl  # noqa: E40214from htr.utils.provenance import run_header  # noqa: E402151617def main() -> None:18    ap = argparse.ArgumentParser()19    ap.add_argument("--pairs", required=True)20    ap.add_argument("--splits", required=True)21    ap.add_argument("--split", default="test")22    ap.add_argument("--checkpoint", default=None)23    ap.add_argument("--system", default=None, help="name for the outputs file (default: checkpoint release name)")24    ap.add_argument("--baselines", action="store_true")25    ap.add_argument("--out", required=True)26    ap.add_argument("--limit", type=int, default=None)27    ap.add_argument("--beam", type=int, default=1)28    ap.add_argument("--device", default="auto")29    ap.add_argument("--variants", default=None)30    ap.add_argument("--batch-size", type=int, default=32, help="batched greedy decoding (1 = item by item)")31    a = ap.parse_args()32    docs = {s["doc_id"] for s in read_jsonl(a.splits) if s["split"] == a.split}33    variants = set(a.variants.split(",")) if a.variants else None34    pairs = [p for p in read_jsonl(a.pairs) if p["doc_id"] in docs and (variants is None or p["variant"] in variants)]35    if a.limit:36        pairs = pairs[: a.limit]37    out = Path(a.out)38    out.mkdir(parents=True, exist_ok=True)39    if a.baselines:40        for name, fn in BASELINES.items():41            write_jsonl(out / f"outputs_{name}.jsonl", ({"pair_id": p["pair_id"], "output": fn(p["source_text"], p["metadata"])} for p in pairs))42            print(name, len(pairs))43    if a.checkpoint:44        from htr.inference.pipeline import HTRReconstructor4546        rec = HTRReconstructor.from_pretrained(a.checkpoint, device=a.device)47        name = a.system or rec.meta.get("name", "htr")48        t0 = time.time()49        rows = []50        if a.beam <= 1 and a.batch_size > 1:51            for b in range(0, len(pairs), 512):52                chunk = pairs[b : b + 512]53                res = rec.reconstruct_many([p["source_text"] for p in chunk], [p["metadata"] for p in chunk], batch_size=a.batch_size)54                rows.extend({"pair_id": p["pair_id"], "output": r.reconstruction, "warnings": r.warnings} for p, r in zip(chunk, res))55                print(f"{len(rows)}/{len(pairs)} {time.time()-t0:.0f}s")56        else:57            for i, p in enumerate(pairs):58                r = rec.reconstruct(p["source_text"], p["metadata"], beam=a.beam)59                rows.append({"pair_id": p["pair_id"], "output": r.reconstruction, "warnings": r.warnings})60                if (i + 1) % 100 == 0:61                    print(f"{i+1}/{len(pairs)} {time.time()-t0:.0f}s")62        write_jsonl(out / f"outputs_{name}.jsonl", rows)63        write_json(out / f"predict_{name}.json", {**run_header(), "checkpoint": a.checkpoint, "n": len(rows), "beam": a.beam, "seconds": round(time.time() - t0, 1), "split": a.split})64        print(name, len(rows), f"{time.time()-t0:.0f}s")656667if __name__ == "__main__":68    main()69