spb/wpai0002_uqo

1#!/usr/bin/env python32"""Edit the pairs of a split with an HTR-Edit checkpoint → outputs_<system>.jsonl (same interface as predict.py)."""3from __future__ import annotations45import argparse6import sys7import time8from pathlib import Path910ROOT = Path(__file__).resolve().parents[1]11sys.path.insert(0, str(ROOT / "src"))12from htr.utils.io import read_jsonl, write_json, write_jsonl  # noqa: E40213from htr.utils.provenance import run_header  # noqa: E402141516def main() -> None:17    ap = argparse.ArgumentParser()18    ap.add_argument("--pairs", required=True)19    ap.add_argument("--splits", required=True)20    ap.add_argument("--split", default="test")21    ap.add_argument("--checkpoint", required=True, help="HTR-Edit checkpoint dir (config.json model_type htr_edit, phrases.json, tokenizer.model)")22    ap.add_argument("--system", default=None, help="name for the outputs file (default: release name)")23    ap.add_argument("--out", required=True)24    ap.add_argument("--limit", type=int, default=None)25    ap.add_argument("--device", default="auto")26    ap.add_argument("--variants", default=None)27    ap.add_argument("--batch-size", type=int, default=64)28    ap.add_argument("--none-bias", type=float, default=None, help="override the <none> logit adjustment stored in config.json")29    a = ap.parse_args()30    docs = {s["doc_id"] for s in read_jsonl(a.splits) if s["split"] == a.split}31    variants = set(a.variants.split(",")) if a.variants else None32    pairs = [p for p in read_jsonl(a.pairs) if p["doc_id"] in docs and (variants is None or p["variant"] in variants)]33    if a.limit:34        pairs = pairs[: a.limit]35    out = Path(a.out)36    out.mkdir(parents=True, exist_ok=True)37    from htr.inference.editor import HTREditor3839    ed = HTREditor.from_pretrained(a.checkpoint, device=a.device, none_bias=a.none_bias)40    name = a.system or ed.meta.get("name", "htr-edit")41    t0 = time.time()42    rows, n_del, n_ins, n_tok = [], 0, 0, 043    for b in range(0, len(pairs), 1024):44        chunk = pairs[b : b + 1024]45        res = ed.edit_many([p["source_text"] for p in chunk], [p["metadata"] for p in chunk], batch_size=a.batch_size)46        for p, r in zip(chunk, res):47            rows.append({"pair_id": p["pair_id"], "output": r.reconstruction, "warnings": r.warnings})48            e = r.metadata_used.get("edits", {})49            n_del += e.get("deleted_tokens", 0)50            n_ins += e.get("insertions", 0)51            n_tok += e.get("source_tokens", 0)52        print(f"{len(rows)}/{len(pairs)} {time.time()-t0:.0f}s", flush=True)53    write_jsonl(out / f"outputs_{name}.jsonl", rows)54    unchanged = sum(1 for p, r in zip(pairs, rows) if r["output"] == p["source_text"])55    write_json(out / f"predict_{name}.json", {**run_header(), "checkpoint": a.checkpoint, "model_type": "htr_edit", "n": len(rows), "seconds": round(time.time() - t0, 1), "split": a.split, "none_bias": ed.none_bias if ed.none_bias is not None else ed.model.cfg.none_logit_bias, "edit_stats": {"deleted_token_rate": n_del / max(1, n_tok), "insertions_per_100_tokens": 100 * n_ins / max(1, n_tok), "outputs_identical_to_source": unchanged}})56    print(name, len(rows), f"{time.time()-t0:.0f}s")575859if __name__ == "__main__":60    main()61