spb/wpai0002_uqo

1#!/usr/bin/env python32"""Predict a split with a fine-tuned B2 (T5-family) run → outputs_B2_<model>.jsonl."""3from __future__ import annotations45import argparse6import sys7from pathlib import Path89ROOT = Path(__file__).resolve().parents[1]10sys.path.insert(0, str(ROOT / "src"))11from htr.training.hf_finetune import predict_hf  # noqa: E40212from htr.utils.io import read_jsonl, write_jsonl  # noqa: E402131415def main() -> None:16    ap = argparse.ArgumentParser()17    ap.add_argument("--pairs", required=True)18    ap.add_argument("--splits", required=True)19    ap.add_argument("--split", default="val")20    ap.add_argument("--run", required=True, help="run dir containing best/")21    ap.add_argument("--name", default="B2_t5-small")22    ap.add_argument("--out", required=True)23    ap.add_argument("--device", default="auto")24    ap.add_argument("--limit", type=int, default=None)25    a = ap.parse_args()26    docs = {s["doc_id"] for s in read_jsonl(a.splits) if s["split"] == a.split}27    pairs = [p for p in read_jsonl(a.pairs) if p["doc_id"] in docs][: a.limit or None]28    outs = predict_hf(Path(a.run), pairs, a.device)29    Path(a.out).mkdir(parents=True, exist_ok=True)30    write_jsonl(Path(a.out) / f"outputs_{a.name}.jsonl", ({"pair_id": k, "output": v} for k, v in outs.items()))31    print(a.name, a.split, len(outs))323334if __name__ == "__main__":35    main()36