Repository Public
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 re8import time9from pathlib import Path1011import torch1213if torch.cuda.is_available() and hasattr(torch.backends.cuda, "enable_cudnn_sdp"):14 torch.backends.cuda.enable_cudnn_sdp(False) # torch 2.11/cu128: cuDNN SDPA gives NaN gradients in bf16 and rebuilds a plan per shape1516from htr.training.data import meta_prefix17from htr.utils.io import append_jsonl, read_jsonl, write_json18from htr.utils.provenance import run_header192021class _null:22 def __enter__(self):23 return None2425 def __exit__(self, *a):26 return False272829def _device(name: str) -> torch.device:30 if name != "auto":31 return torch.device(name)32 return torch.device("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"))333435_PH = re.compile(r"⟦([CMU])(\d+)⟧")36_PH_ASCII = re.compile(r"\[\[([CMU])(\d+)\]\]")373839def ph_to_ascii(text: str) -> str:40 """⟦C1⟧ → [[C1]]: the pretrained tokenizers map ⟦/⟧ to <unk>, and resizing the embeddings for added tokens breaks41 T5 v1.1 (initial CE ≈ 800 instead of 1.3 with transformers 5.22 / torch 2.11), so placeholders are spelled in ASCII."""42 return _PH.sub(lambda m: f"[[{m.group(1)}{m.group(2)}]]", text)434445def ascii_to_ph(text: str) -> str:46 return _PH_ASCII.sub(lambda m: f"⟦{m.group(1)}{m.group(2)}⟧", text)474849def protect_unk_chars(text: str, tok) -> tuple[str, dict[str, str]]:50 """T5's SentencePiece vocabulary has no piece for '<', '~', '{', '}', '^', '\\', '±', '→', '≤', Greek letters, … (they become51 <unk> and are silently dropped: 'p < 0.01' → 'p 0.01'). Replace each distinct unknown character by a bracketed ASCII52 placeholder [[S1]], [[S2]], … that the fine-tuned models copy like the citation placeholders, and restore after decoding."""53 unk = tok.unk_token_id54 cache: dict[str, bool] = {}55 mapping: dict[str, str] = {}56 out = []57 for ch in text:58 if ch.isspace() or (ch.isascii() and ch.isalnum()):59 out.append(ch)60 continue61 if ch not in cache:62 cache[ch] = unk in tok(ch, add_special_tokens=False).input_ids63 if cache[ch]:64 if ch not in mapping:65 mapping[ch] = f"[[S{len(mapping) + 1}]]"66 out.append(mapping[ch])67 else:68 out.append(ch)69 return "".join(out), mapping707172def restore_unk_chars(text: str, mapping: dict[str, str]) -> str:73 for ch, ph in mapping.items():74 text = text.replace(ph, ch)75 return text767778def load_hf(model_name: str, device: torch.device):79 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer8081 tok = AutoTokenizer.from_pretrained(model_name)82 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)83 return tok, model.to(device)848586def 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, patience: int = 6, warmup: int = 200, grad_ckpt: bool = False) -> dict:87 torch.manual_seed(seed)88 dev = _device(device)89 autocast = torch.autocast("cuda", dtype=torch.bfloat16) if dev.type == "cuda" else None # bf16 on GPU boxes (BART/T5-large)90 tok, model = load_hf(model_name, dev)91 if grad_ckpt:92 model.gradient_checkpointing_enable()93 model.config.use_cache = False94 splits = {s["doc_id"]: s["split"] for s in read_jsonl(splits_path)}95 pairs = [p for p in read_jsonl(pairs_path) if variants is None or p["variant"] in variants]96 tr = [p for p in pairs if splits.get(p["doc_id"]) == "train"][: limit or None]97 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 slow98 run_dir.mkdir(parents=True, exist_ok=True)99 opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)100101 def batches(rows, shuffle):102 idx = torch.randperm(len(rows)).tolist() if shuffle else list(range(len(rows)))103 for i in range(0, len(idx), batch_size):104 chunk = [rows[j] for j in idx[i : i + batch_size]]105 enc = tok(["reconstruct: " + meta_prefix(p) + ph_to_ascii(p["source_text"]) for p in chunk], return_tensors="pt", padding=True, pad_to_multiple_of=32, truncation=True, max_length=max_len)106 lab = tok([ph_to_ascii(p["target_text"]) for p in chunk], return_tensors="pt", padding=True, pad_to_multiple_of=32, truncation=True, max_length=max_len).input_ids107 lab[lab == tok.pad_token_id] = -100108 yield {k: v.to(dev) for k, v in enc.items()}, lab.to(dev)109110 @torch.no_grad()111 def evaluate():112 model.eval()113 tot, n = 0.0, 0114 for enc, lab in batches(va, False):115 with (autocast or _null()):116 out = model(**enc, labels=lab)117 k = int((lab != -100).sum())118 tot += float(out.loss) * k119 n += k120 model.train()121 return tot / max(1, n)122123 step, best, t0, accel, bad = 0, math.inf, time.time(), 0.0, 0124 model.train()125 while step < max_steps and bad < patience:126 for enc, lab in batches(tr, True):127 t1 = time.perf_counter()128 for g in opt.param_groups:129 g["lr"] = lr * min(1.0, (step + 1) / warmup) * (0.5 * (1 + math.cos(math.pi * min(1.0, step / max_steps))) if step >= warmup else 1.0)130 with (autocast or _null()):131 loss = model(**enc, labels=lab).loss132 loss.backward()133 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)134 opt.step()135 opt.zero_grad(set_to_none=True)136 if dev.type == "mps":137 torch.mps.synchronize()138 if step % 20 == 0:139 torch.mps.empty_cache() # HF models + dynamic shapes leak MPS allocations otherwise140 accel += time.perf_counter() - t1141 step += 1142 if step % 25 == 0:143 append_jsonl(run_dir / "train_log.jsonl", {"step": step, "ce": float(loss.detach()), "elapsed_s": round(time.time() - t0, 1)})144 if step % eval_every == 0 or step >= max_steps:145 ce = evaluate()146 append_jsonl(run_dir / "train_log.jsonl", {"step": step, "eval": {"val_ce": ce}})147 if ce < best:148 best, bad = ce, 0149 model.save_pretrained(run_dir / "best")150 tok.save_pretrained(run_dir / "best")151 else:152 bad += 1153 if step >= max_steps or bad >= patience:154 break155 n_params = sum(p.numel() for p in model.parameters())156 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)}157 write_json(run_dir / "train_summary.json", summary)158 return summary159160161@torch.no_grad()162def 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]:163 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer164165 dev = _device(device)166 tok = AutoTokenizer.from_pretrained(run_dir / "best")167 model = AutoModelForSeq2SeqLM.from_pretrained(run_dir / "best").to(dev).eval()168 out: dict[str, str] = {}169 for i in range(0, len(pairs), batch_size):170 chunk = pairs[i : i + batch_size]171 prot = [protect_unk_chars(ph_to_ascii(p["source_text"]), tok) for p in chunk]172 enc = tok(["reconstruct: " + meta_prefix(p) + pr[0] for p, pr in zip(chunk, prot)], return_tensors="pt", padding=True, truncation=True, max_length=max_len).to(dev)173 with (torch.autocast("cuda", dtype=torch.bfloat16) if dev.type == "cuda" else _null()):174 gen = model.generate(**enc, max_new_tokens=max_len, num_beams=beam, do_sample=False)175 for p, pr, g in zip(chunk, prot, gen):176 out[p["pair_id"]] = ascii_to_ph(restore_unk_chars(tok.decode(g, skip_special_tokens=True), pr[1]))177 return out178