spb/wpai0002_uqo

1#!/usr/bin/env python32"""B2 baseline: fine-tune a pretrained T5-family model (downloads weights on first use; opt-in)."""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 finetune, 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("--run", required=True)20    ap.add_argument("--model-name", default="google-t5/t5-small")21    ap.add_argument("--max-steps", type=int, default=3000)22    ap.add_argument("--batch-size", type=int, default=16)23    ap.add_argument("--lr", type=float, default=3e-4)24    ap.add_argument("--device", default="auto")25    ap.add_argument("--limit", type=int, default=None)26    ap.add_argument("--variants", default=None)27    ap.add_argument("--predict-split", default=None, help="after training, predict this split into --pred-out")28    ap.add_argument("--pred-out", default=None)29    ap.add_argument("--pred-limit", type=int, default=None)30    ap.add_argument("--patience", type=int, default=6)31    ap.add_argument("--max-len", type=int, default=512)32    ap.add_argument("--grad-ckpt", action="store_true", help="gradient checkpointing (flan-t5-xl on one 80 GB GPU)")33    ap.add_argument("--warmup", type=int, default=200)34    ap.add_argument("--eval-every", type=int, default=250)35    ap.add_argument("--system", default=None, help="outputs file name (default B2_<model>)")36    ap.add_argument("--pred-splits", default=None, help="comma-separated extra splits to predict into --pred-out")37    a = ap.parse_args()38    run_dir = ROOT / "artifacts/runs" / a.run39    variants = set(a.variants.split(",")) if a.variants else None40    s = finetune(Path(a.pairs), Path(a.splits), run_dir, a.model_name, a.max_steps, a.lr, a.batch_size, device=a.device, limit=a.limit, variants=variants, patience=a.patience, warmup=a.warmup, eval_every=a.eval_every, max_len=a.max_len, grad_ckpt=a.grad_ckpt)41    print("summary:", s)42    system = a.system or f"B2_{a.model_name.split('/')[-1]}"43    todo = ([a.predict_split] if a.predict_split else []) + (a.pred_splits.split(",") if a.pred_splits else [])44    for spl in todo:45        docs = {x["doc_id"] for x in read_jsonl(a.splits) if x["split"] == spl}46        pairs = [p for p in read_jsonl(a.pairs) if p["doc_id"] in docs and (variants is None or p["variant"] in variants)]47        if spl == "val" and a.pred_limit:48            pairs = pairs[: a.pred_limit]49        outs = predict_hf(run_dir, pairs, a.device, batch_size=32)50        out_dir = Path(a.pred_out) / f"medium_{spl}" if a.pred_out and not a.pred_out.endswith(f"medium_{spl}") else Path(a.pred_out)51        out_dir.mkdir(parents=True, exist_ok=True)52        write_jsonl(out_dir / f"outputs_{system}.jsonl", ({"pair_id": p["pair_id"], "doc_id": p["doc_id"], "system": system, "output": outs[p["pair_id"]]} for p in pairs))53        print("predicted", spl, len(outs), "->", out_dir)545556if __name__ == "__main__":57    main()58