spb/wpai0002_uqo

1"""``htr`` command-line interface (versioned): reconstruct text, batch files, inspect a release."""2from __future__ import annotations34import argparse5import json6import sys7from pathlib import Path89from htr import __version__101112def main(argv: list[str] | None = None) -> int:13    ap = argparse.ArgumentParser(prog="htr", description="Human Text Reconstruction Transformer (research tool; outputs are synthetic reconstructions)")14    ap.add_argument("--version", action="version", version=f"htr {__version__}")15    sub = ap.add_subparsers(dest="cmd", required=True)16    r = sub.add_parser("reconstruct", help="reconstruct one passage")17    r.add_argument("--model", required=True)18    r.add_argument("--text", default=None)19    r.add_argument("--file", default=None)20    r.add_argument("--discipline", default=None)21    r.add_argument("--year", type=int, default=None)22    r.add_argument("--role", default=None)23    r.add_argument("--beam", type=int, default=1)24    r.add_argument("--device", default="auto")25    r.add_argument("--json", action="store_true")26    b = sub.add_parser("batch", help="reconstruct a JSONL file with fields text[, discipline, pub_year, section_role]")27    b.add_argument("--model", required=True)28    b.add_argument("--input", required=True)29    b.add_argument("--output", required=True)30    b.add_argument("--beam", type=int, default=1)31    b.add_argument("--device", default="auto")32    i = sub.add_parser("info", help="show release metadata")33    i.add_argument("--model", required=True)34    a = ap.parse_args(argv)35    if a.cmd == "info":36        p = Path(a.model)37        for f in ("release_meta.json", "config.json", "MANIFEST.json"):38            if (p / f).exists():39                print(f"== {f}")40                print((p / f).read_text()[:2000])41        return 042    from htr.inference.pipeline import HTRReconstructor4344    rec = HTRReconstructor.from_pretrained(a.model, device=a.device)45    if a.cmd == "reconstruct":46        text = a.text or (Path(a.file).read_text() if a.file else sys.stdin.read())47        res = rec.reconstruct(text.strip(), {"discipline": a.discipline, "pub_year": a.year, "section_role": a.role}, beam=a.beam)48        if a.json:49            print(json.dumps(res.to_dict(), ensure_ascii=False, indent=2))50        else:51            print(res.reconstruction)52            for w in res.warnings:53                print(f"! {w}", file=sys.stderr)54            print(f"[{res.notice}]", file=sys.stderr)55        return 056    if a.cmd == "batch":57        n = 058        with open(a.input, encoding="utf-8") as fin, open(a.output, "w", encoding="utf-8") as fout:59            for line in fin:60                if not line.strip():61                    continue62                row = json.loads(line)63                res = rec.reconstruct(row["text"], {k: row.get(k) for k in ("discipline", "pub_year", "section_role", "era")}, beam=a.beam)64                fout.write(json.dumps({**row, "reconstruction": res.reconstruction, "warnings": res.warnings, "preservation": res.preservation}, ensure_ascii=False) + "\n")65                n += 166        print(f"wrote {n} reconstructions to {a.output}", file=sys.stderr)67        return 068    return 1697071if __name__ == "__main__":72    raise SystemExit(main())73