Repository Public Research preview
1#!/usr/bin/env python32"""Train HTR-Edit (non-generative KEEP/DELETE + phrase-insertion tagger) on validated pairs.34Steps: tokenise + align train/val pairs → phrase vocabulary (train only) → edit labels + coverage statistics →5oracle (gold-label) reconstruction scores on validation → encoder initialised from a denoiser/HTR checkpoint →6AdamW, warmup + cosine, gradient clipping, evaluation every N steps (validation loss, token accuracies, exact7match), early stopping on validation loss, SafeTensors checkpoints (``best/``, ``last/``) and a manifest.89Memory caps on MPS nodes are set by the caller:10``PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.7 PYTORCH_MPS_LOW_WATERMARK_RATIO=0.5``.11"""12from __future__ import annotations1314import argparse15import math16import os17import shutil18import sys19import time20from pathlib import Path2122ROOT = Path(__file__).resolve().parents[1]23sys.path.insert(0, str(ROOT / "src"))24import numpy as np # noqa: E40225import torch # noqa: E4022627from htr.models.tagger import EditConfig, HTREditModel, init_from_checkpoint, save_edit_checkpoint # noqa: E40228from htr.models.tokenizer import HTRTokenizer # noqa: E40229from htr.training.data import PAD_MULTIPLE, BucketSampler # noqa: E40230from htr.training.edits import PhraseVocab, apply_edits, coverage_stats, labels_from_spans, protected_ids, protection_plan, space_ids, tokenize_align_many # noqa: E40231from htr.utils.io import append_jsonl, read_jsonl, write_json # noqa: E40232from htr.utils.provenance import run_header # noqa: E402333435def device_of(name: str) -> torch.device:36 if name != "auto":37 return torch.device(name)38 if torch.cuda.is_available():39 return torch.device("cuda")40 return torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")414243class EditData:44 """Pre-computed (ids, tags, inserts) arrays; ``target`` ids kept for exact-match evaluation."""4546 def __init__(self, items: list[tuple[list[int], list[int], list[int], list[int]]]):47 self.src = [np.asarray(s, dtype=np.int32) for s, _, _, _ in items]48 self.tags = [np.asarray(t, dtype=np.int8) for _, t, _, _ in items]49 self.ins = [np.asarray(i, dtype=np.int16) for _, _, i, _ in items]50 self.tgt = [np.asarray(g, dtype=np.int32) for _, _, _, g in items]5152 def __len__(self) -> int:53 return len(self.src)5455 def batch(self, idx: list[int], cfg: EditConfig) -> dict:56 n_max = max(len(self.src[i]) for i in idx)57 L = n_max + 258 L = ((L + PAD_MULTIPLE - 1) // PAD_MULTIPLE) * PAD_MULTIPLE59 ids = np.full((len(idx), L), cfg.pad_id, dtype=np.int64)60 tags = np.full((len(idx), L), -100, dtype=np.int64)61 ins = np.full((len(idx), L - 1), -100, dtype=np.int64)62 for r, i in enumerate(idx):63 n = len(self.src[i])64 ids[r, 0], ids[r, n + 1] = cfg.bos_id, cfg.eos_id65 ids[r, 1 : n + 1] = self.src[i]66 tags[r, 1 : n + 1] = self.tags[i]67 ins[r, : n + 1] = self.ins[i]68 return {"ids": torch.from_numpy(ids), "tags": torch.from_numpy(tags), "ins": torch.from_numpy(ins), "idx": idx}697071def lr_at(step: int, base: float, warmup: int, max_steps: int) -> float:72 if step < warmup:73 return base * (step + 1) / warmup74 prog = min(1.0, (step - warmup) / max(1, max_steps - warmup))75 return base * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * prog)))767778@torch.no_grad()79def evaluate(model: HTREditModel, data: EditData, batches: list[list[int]], dev: torch.device, phrases: PhraseVocab, keep_ids: set[int], spaced: set[int]) -> dict:80 model.eval()81 tot = w = 0.082 tag_ok = tag_n = ins_ok = ins_n = edit_ok = edit_n = del_ok = del_n = del_pred = em_lab = em_tgt = n = 083 for idx in batches:84 b = data.batch(idx, model.cfg)85 bd = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in b.items()}86 loss, logs = model.loss(bd)87 tot += float(loss) * len(idx)88 w += len(idx)89 pt, pi = model.predict(bd["ids"])90 pt, pi = pt.cpu().numpy(), pi.cpu().numpy()91 gt, gi = b["tags"].numpy(), b["ins"].numpy()92 mt, mi = gt != -100, gi != -10093 tag_ok += int(((pt == gt) & mt).sum())94 tag_n += int(mt.sum())95 ins_ok += int(((pi == gi) & mi).sum())96 ins_n += int(mi.sum())97 edit_ok += int(((pi == gi) & (gi > 0)).sum()) # insertion recall (exact phrase on gold insertion slots)98 edit_n += int((gi > 0).sum())99 del_ok += int(((pt == 1) & (gt == 1)).sum()) # deletion recall100 del_n += int((gt == 1).sum())101 del_pred += int(((pt == 1) & mt).sum())102 for r, i in enumerate(idx):103 k = len(data.src[i])104 same = bool((pt[r, 1 : k + 1] == gt[r, 1 : k + 1]).all() and (pi[r, : k + 1] == gi[r, : k + 1]).all())105 em_lab += same106 srcl = data.src[i].tolist()107 keep, block = protection_plan(srcl, keep_ids, spaced)108 out = apply_edits(srcl, pt[r, 1 : k + 1].tolist(), pi[r, : k + 1].tolist(), phrases, force_keep=keep, block_slot=block)109 em_tgt += out == data.tgt[i].tolist()110 n += 1111 model.train()112 return {"val_loss": tot / max(1, w), "tag_acc": tag_ok / max(1, tag_n), "ins_acc": ins_ok / max(1, ins_n), "insert_recall": edit_ok / max(1, edit_n), "delete_recall": del_ok / max(1, del_n), "delete_precision": del_ok / max(1, del_pred), "exact_match_labels": em_lab / max(1, n), "exact_match_target": em_tgt / max(1, n), "n_pairs": n}113114115def main() -> None:116 ap = argparse.ArgumentParser()117 ap.add_argument("--pairs", required=True)118 ap.add_argument("--splits", required=True)119 ap.add_argument("--tokenizer", required=True, help="dir with tokenizer.model")120 ap.add_argument("--init-from", default=None, help="HTR checkpoint dir whose emb.*/enc.*/enc_ln.* tensors initialise the encoder")121 ap.add_argument("--run", required=True)122 ap.add_argument("--size", default="tiny")123 ap.add_argument("--phrases", type=int, default=4000, help="V_phr (plus <none>)")124 ap.add_argument("--phrase-max-len", type=int, default=4)125 ap.add_argument("--max-src", type=int, default=510)126 ap.add_argument("--max-steps", type=int, default=15000)127 ap.add_argument("--eval-every", type=int, default=1000)128 ap.add_argument("--log-every", type=int, default=50)129 ap.add_argument("--lr", type=float, default=3e-4)130 ap.add_argument("--warmup", type=int, default=500)131 ap.add_argument("--weight-decay", type=float, default=0.01)132 ap.add_argument("--max-tokens", type=int, default=4096)133 ap.add_argument("--max-batch", type=int, default=64)134 ap.add_argument("--patience", type=int, default=4)135 ap.add_argument("--val-limit", type=int, default=2000, help="validation pairs used for model selection")136 ap.add_argument("--oracle-n", type=int, default=3000, help="validation pairs for the gold-label (oracle) reconstruction score")137 ap.add_argument("--seed", type=int, default=1)138 ap.add_argument("--device", default="auto")139 ap.add_argument("--workers", type=int, default=8)140 ap.add_argument("--limit", type=int, default=None, help="cap on training pairs (smoke tests)")141 a = ap.parse_args()142 torch.manual_seed(a.seed)143 np.random.seed(a.seed)144 t_all = time.time()145 run_dir = ROOT / "artifacts/runs" / a.run146 run_dir.mkdir(parents=True, exist_ok=True)147 tok_path = Path(a.tokenizer) / "tokenizer.model"148 tok = HTRTokenizer(tok_path)149 prot = protected_ids(tok.sp)150 spaced = space_ids(tok.sp)151152 splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}153 tr_pairs, va_pairs = [], []154 for p in read_jsonl(a.pairs):155 s = splits.get(p["doc_id"])156 if s == "train":157 tr_pairs.append(p)158 elif s == "val":159 va_pairs.append(p)160 if a.limit:161 tr_pairs = tr_pairs[: a.limit]162 n_val_keep = max(a.val_limit, a.oracle_n)163 va_pairs = va_pairs[:n_val_keep]164 print(f"pairs: train {len(tr_pairs)} · val (kept) {len(va_pairs)}", flush=True)165166 t0 = time.time()167 tr_al = tokenize_align_many([(p["source_text"], p["target_text"]) for p in tr_pairs], str(tok_path), a.workers)168 va_al = tokenize_align_many([(p["source_text"], p["target_text"]) for p in va_pairs], str(tok_path), a.workers)169 print(f"tokenised + aligned in {time.time()-t0:.0f}s", flush=True)170171 phrases = PhraseVocab.build((s for _, _, s in tr_al), size=a.phrases, max_len=a.phrase_max_len, protected=prot)172 phrases.save(run_dir / "phrases.json", tok.sp)173 print(f"phrase vocabulary: {len(phrases)} (incl. <none>); meta {phrases.meta}", flush=True)174175 def build(al):176 labs, items, skipped = [], [], 0177 for src, tgt, spans in al:178 lab = labels_from_spans(src, tgt, spans, phrases)179 labs.append(lab)180 if len(src) > a.max_src or len(src) == 0:181 skipped += 1182 items.append(None)183 continue184 items.append((src, lab.tags, lab.inserts, tgt))185 return labs, items, skipped186187 tr_labs, tr_items, tr_skip = build(tr_al)188 va_labs, va_items, va_skip = build(va_al)189 cov = {"train": {**coverage_stats(tr_labs), "skipped_too_long": tr_skip}, "val": {**coverage_stats(va_labs), "skipped_too_long": va_skip}, "max_src": a.max_src}190191 # oracle: realise the gold labels on the first oracle_n validation pairs (= the predict_tagger --limit subset)192 from htr.evaluation.metrics import lexical193194 orc, ident = [], []195 for p, (src, tgt, _), lab in list(zip(va_pairs, va_al, va_labs))[: a.oracle_n]:196 hyp = tok.sp.decode(apply_edits(src, lab.tags, lab.inserts, phrases))197 orc.append(lexical(hyp, p["target_text"])["chrf"])198 ident.append(lexical(tok.sp.decode(src), p["target_text"])["chrf"])199 cov["oracle_val"] = {"n": len(orc), "chrf_oracle_labels_mean": float(np.mean(orc)) if orc else None, "chrf_identity_after_tokenizer_roundtrip_mean": float(np.mean(ident)) if ident else None}200 write_json(run_dir / "coverage.json", {**run_header(), **cov})201 print("coverage:", cov, flush=True)202 del tr_al, va_al, tr_labs, va_labs203204 tr = EditData([x for x in tr_items if x is not None])205 va_all = [x for x in va_items if x is not None]206 va = EditData(va_all[: a.val_limit])207 del tr_items, va_items208209 cfg = EditConfig.from_size(a.size, tok.vocab_size, max_src=a.max_src, n_phrases=len(phrases))210 model = HTREditModel(cfg)211 init_info = init_from_checkpoint(model, a.init_from) if a.init_from else None212 if init_info:213 print(f"initialised {init_info['loaded']}/{init_info['encoder_tensors']} encoder tensors from {a.init_from}; new: {init_info['new_tensors']}", flush=True)214 dev = device_of(a.device)215 model.to(dev).train()216 print(f"HTR-Edit/{a.size}: {model.n_params/1e6:.1f}M params on {dev}; train {len(tr)} val {len(va)}", flush=True)217218 shutil.copy(tok_path, run_dir / "tokenizer.model")219 write_json(run_dir / "manifest.json", {**run_header({"run": a.run}), "model_type": "htr_edit", "args": vars(a), "env": {k: os.environ.get(k) for k in ("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "PYTORCH_MPS_LOW_WATERMARK_RATIO")}, "init": init_info, "model_config": cfg.to_dict(), "n_params": model.n_params, "n_train_pairs": len(tr), "n_val_pairs_selection": len(va), "phrase_vocab": phrases.meta, "coverage": cov, "loss": {"tag": "CE", "insert": f"CE weighted {cfg.insert_w_edit} on insertion slots, {cfg.insert_w_none} on <none>", "none_logit_bias_at_inference": cfg.none_logit_bias}})220221 decay, no_decay = [], []222 for n_, p_ in model.named_parameters():223 (no_decay if p_.ndim < 2 or "ln" in n_ else decay).append(p_)224 opt = torch.optim.AdamW([{"params": decay, "weight_decay": a.weight_decay}, {"params": no_decay, "weight_decay": 0.0}], lr=a.lr, betas=(0.9, 0.98), eps=1e-8)225 sampler = BucketSampler([len(s) + 2 for s in tr.src], a.max_tokens, a.max_batch, a.seed, True)226 val_batches = BucketSampler([len(s) + 2 for s in va.src], a.max_tokens, a.max_batch, 0, False)._batches()227 log = run_dir / "train_log.jsonl"228 if log.exists():229 log.unlink()230 step, epoch, best, bad, tokens = 0, 0, float("inf"), 0, 0231 best_eval, t_start, done = None, time.time(), False232 while not done:233 sampler.set_epoch(epoch)234 for idx in sampler:235 b = tr.batch(idx, cfg)236 bd = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in b.items()}237 loss, logs = model.loss(bd)238 loss.backward()239 gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)240 for g in opt.param_groups:241 g["lr"] = lr_at(step, a.lr, a.warmup, a.max_steps)242 opt.step()243 opt.zero_grad(set_to_none=True)244 tokens += logs["n_tok"]245 step += 1246 if dev.type == "mps" and step % 20 == 0:247 torch.mps.empty_cache()248 if step % a.log_every == 0:249 mem = round(torch.mps.driver_allocated_memory() / 2**30, 2) if dev.type == "mps" else None250 append_jsonl(log, {"step": step, "epoch": epoch, "lr": lr_at(step, a.lr, a.warmup, a.max_steps), "grad_norm": float(gn), "tokens_seen": tokens, "elapsed_s": round(time.time() - t_start, 1), "mem_gb": mem, **logs})251 if not math.isfinite(logs["total"]):252 raise RuntimeError(f"non-finite loss at step {step}: {logs}")253 if step % a.eval_every == 0 or step >= a.max_steps:254 ev = evaluate(model, va, val_batches, dev, phrases, prot, spaced)255 append_jsonl(log, {"step": step, "eval": ev})256 print(f"step {step} ep {epoch} {time.time()-t_start:.0f}s train {logs['total']:.3f} | val {ev}", flush=True)257 extra = {"step": step, "eval": ev, "tokens_seen": tokens}258 if ev["val_loss"] < best - 1e-4:259 best, bad, best_eval = ev["val_loss"], 0, {"step": step, **ev}260 save_edit_checkpoint(model, run_dir / "best", extra)261 else:262 bad += 1263 save_edit_checkpoint(model, run_dir / "last", extra)264 if dev.type == "mps":265 torch.mps.empty_cache()266 if bad >= a.patience or step >= a.max_steps:267 done = True268 break269 epoch += 1270 for tag in ("best", "last"):271 d = run_dir / tag272 if d.exists():273 shutil.copy(tok_path, d / "tokenizer.model")274 shutil.copy(run_dir / "phrases.json", d / "phrases.json")275 write_json(d / "release_meta.json", {"name": f"htr-edit-{a.size}", "model_type": "htr_edit", "run": a.run})276 summary = {"steps": step, "epochs": epoch, "best": best_eval, "tokens_seen": tokens, "wall_seconds": round(time.time() - t_start, 1), "total_seconds_incl_preprocessing": round(time.time() - t_all, 1), "n_params": model.n_params, "device": str(dev), "stopped_early": step < a.max_steps}277 write_json(run_dir / "train_summary.json", {**run_header(), **summary})278 print("summary:", summary, flush=True)279280281if __name__ == "__main__":282 main()283