spb/wpai0002_uqo

1"""B2 baseline: fine-tune a pretrained encoder–decoder (T5 family, Apache-2.0) on the same pairs with the same2splits. Weights are downloaded from the Hugging Face Hub under their own licence on first use (network; opt-in).3Metadata is given as a text prefix. Outputs go through the same placeholder protocol as HTR models."""4from __future__ import annotations56import math7import time8from pathlib import Path910import torch1112from htr.training.data import meta_prefix13from htr.utils.io import append_jsonl, read_jsonl, write_json14from htr.utils.provenance import run_header151617def _device(name: str) -> torch.device:18    if name != "auto":19        return torch.device(name)20    return torch.device("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"))212223def load_hf(model_name: str, device: torch.device):24    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer2526    tok = AutoTokenizer.from_pretrained(model_name)27    # make placeholders atomic tokens28    new = [f"⟦{k}{i}⟧" for k in "CMU" for i in range(1, 33)]29    tok.add_tokens(new)30    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)31    model.resize_token_embeddings(len(tok))32    return tok, model.to(device)333435def finetune(pairs_path: Path, splits_path: Path, run_dir: Path, model_name: str = "google-t5/t5-small", max_steps: int = 3000, lr: float = 3e-4, batch_size: int = 16, max_len: int = 512, eval_every: int = 250, device: str = "auto", seed: int = 1, limit: int | None = None, variants: set[str] | None = None) -> dict:36    torch.manual_seed(seed)37    dev = _device(device)38    tok, model = load_hf(model_name, dev)39    splits = {s["doc_id"]: s["split"] for s in read_jsonl(splits_path)}40    pairs = [p for p in read_jsonl(pairs_path) if variants is None or p["variant"] in variants]41    tr = [p for p in pairs if splits.get(p["doc_id"]) == "train"][: limit or None]42    va = [p for p in pairs if splits.get(p["doc_id"]) == "val"][: (limit // 5 if limit else 1500)]  # cap: evaluation of T5 on MPS is slow43    run_dir.mkdir(parents=True, exist_ok=True)44    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)4546    def batches(rows, shuffle):47        idx = torch.randperm(len(rows)).tolist() if shuffle else list(range(len(rows)))48        for i in range(0, len(idx), batch_size):49            chunk = [rows[j] for j in idx[i : i + batch_size]]50            enc = tok(["reconstruct: " + meta_prefix(p) + p["source_text"] for p in chunk], return_tensors="pt", padding=True, pad_to_multiple_of=32, truncation=True, max_length=max_len)51            lab = tok([p["target_text"] for p in chunk], return_tensors="pt", padding=True, pad_to_multiple_of=32, truncation=True, max_length=max_len).input_ids52            lab[lab == tok.pad_token_id] = -10053            yield {k: v.to(dev) for k, v in enc.items()}, lab.to(dev)5455    @torch.no_grad()56    def evaluate():57        model.eval()58        tot, n = 0.0, 059        for enc, lab in batches(va, False):60            out = model(**enc, labels=lab)61            k = int((lab != -100).sum())62            tot += float(out.loss) * k63            n += k64        model.train()65        return tot / max(1, n)6667    step, best, t0, accel = 0, math.inf, time.time(), 0.068    model.train()69    while step < max_steps:70        for enc, lab in batches(tr, True):71            t1 = time.perf_counter()72            for g in opt.param_groups:73                g["lr"] = lr * min(1.0, (step + 1) / 200)74            loss = model(**enc, labels=lab).loss75            loss.backward()76            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)77            opt.step()78            opt.zero_grad(set_to_none=True)79            if dev.type == "mps":80                torch.mps.synchronize()81                if step % 20 == 0:82                    torch.mps.empty_cache()  # HF models + dynamic shapes leak MPS allocations otherwise83            accel += time.perf_counter() - t184            step += 185            if step % 25 == 0:86                append_jsonl(run_dir / "train_log.jsonl", {"step": step, "ce": float(loss.detach()), "elapsed_s": round(time.time() - t0, 1)})87            if step % eval_every == 0 or step >= max_steps:88                ce = evaluate()89                append_jsonl(run_dir / "train_log.jsonl", {"step": step, "eval": {"val_ce": ce}})90                if ce < best:91                    best = ce92                    model.save_pretrained(run_dir / "best")93                    tok.save_pretrained(run_dir / "best")94            if step >= max_steps:95                break96    n_params = sum(p.numel() for p in model.parameters())97    summary = {**run_header({"run": run_dir.name}), "model_name": model_name, "steps": step, "best_val_ce": best, "n_params": n_params, "n_train_pairs": len(tr), "n_val_pairs": len(va), "accel_seconds": round(accel, 1), "wall_seconds": round(time.time() - t0, 1), "device": str(dev)}98    write_json(run_dir / "train_summary.json", summary)99    return summary100101102@torch.no_grad()103def predict_hf(run_dir: Path, pairs: list[dict], device: str = "auto", max_len: int = 512, beam: int = 1, batch_size: int = 16) -> dict[str, str]:104    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer105106    dev = _device(device)107    tok = AutoTokenizer.from_pretrained(run_dir / "best")108    model = AutoModelForSeq2SeqLM.from_pretrained(run_dir / "best").to(dev).eval()109    out: dict[str, str] = {}110    for i in range(0, len(pairs), batch_size):111        chunk = pairs[i : i + batch_size]112        enc = tok(["reconstruct: " + meta_prefix(p) + p["source_text"] for p in chunk], return_tensors="pt", padding=True, truncation=True, max_length=max_len).to(dev)113        gen = model.generate(**enc, max_new_tokens=max_len, num_beams=beam, do_sample=False)114        for p, g in zip(chunk, gen):115            out[p["pair_id"]] = tok.decode(g, skip_special_tokens=True)116    return out117