Repository Public Research preview
1#!/usr/bin/env python32"""Train an HTR model (M1/M2/M3 or ablations) on validated pairs with leakage-safe splits."""3from __future__ import annotations45import argparse6import json7import shutil8import sys9from pathlib import Path1011ROOT = Path(__file__).resolve().parents[1]12sys.path.insert(0, str(ROOT / "src"))13import torch # noqa: E4021415from htr.models.config import make_config # noqa: E40216from htr.models.losses import LossWeights # noqa: E40217from htr.models.tokenizer import HTRTokenizer # noqa: E40218from htr.models.transformer import HTRModel # noqa: E40219from htr.training.data import PairDataset, StyleLabelMap, collate, make_loader # noqa: E40220from htr.training.trainer import TrainConfig, Trainer, overfit_one_batch # noqa: E40221from htr.utils.io import read_jsonl, write_json # noqa: E40222from htr.utils.provenance import run_header # noqa: E4022324PRESETS = {25 "m1": dict(arch="m1", loss=dict()),26 "m2": dict(arch="m2", loss=dict()),27 "m3": dict(arch="m3", loss=dict(style=0.3, adv=0.3, con=0.2, sem=0.1, orth=0.1, fact=0.5)),28 "m3_tokenonly": dict(arch="m3", loss=dict()),29 "m3_nostyle": dict(arch="m3", loss=dict(con=0.2, sem=0.1, fact=0.5)),30 "m3_noadv": dict(arch="m3", loss=dict(style=0.3, con=0.2, sem=0.1, orth=0.1, fact=0.5)),31 "m3_nofact": dict(arch="m3", loss=dict(style=0.3, adv=0.3, con=0.2, sem=0.1, orth=0.1)),32 "m3_nocon": dict(arch="m3", loss=dict(style=0.3, adv=0.3, orth=0.1, fact=0.5)),33 # single-term diagnostics (pilot: every M3 variant with auxiliary losses generalised far worse than token-only)34 "m1_fact": dict(arch="m1", loss=dict(fact=0.5)),35 "m3_factonly": dict(arch="m3", loss=dict(fact=0.5)),36 "m3_cononly": dict(arch="m3", loss=dict(con=0.2)),37 "m3_semonly": dict(arch="m3", loss=dict(sem=0.1)),38 # light dual-latent objective for fine-tuning from a denoising-pretrained backbone (D19 follow-up)39 "m3_light": dict(arch="m3", loss=dict(style=0.05, adv=0.05, con=0.1, sem=0.1, orth=0.05, fact=0.2)),40}414243def main() -> None:44 ap = argparse.ArgumentParser()45 ap.add_argument("--pairs", required=True)46 ap.add_argument("--splits", required=True)47 ap.add_argument("--tokenizer", required=True, help="dir with tokenizer.model")48 ap.add_argument("--run", required=True, help="run name under artifacts/runs/")49 ap.add_argument("--preset", default="m1", choices=sorted(PRESETS))50 ap.add_argument("--size", default="tiny")51 ap.add_argument("--max-steps", type=int, default=5000)52 ap.add_argument("--eval-every", type=int, default=250)53 ap.add_argument("--lr", type=float, default=3e-4)54 ap.add_argument("--warmup", type=int, default=300)55 ap.add_argument("--max-tokens", type=int, default=8192)56 ap.add_argument("--max-batch", type=int, default=48)57 ap.add_argument("--accum", type=int, default=1)58 ap.add_argument("--patience", type=int, default=6)59 ap.add_argument("--seed", type=int, default=1)60 ap.add_argument("--device", default="auto")61 ap.add_argument("--num-workers", type=int, default=0, help="DataLoader workers (0 = main process, as on MPS)")62 ap.add_argument("--amp", default="auto")63 ap.add_argument("--variants", default=None, help="restrict training variants, e.g. A,B,C,D (hold out E)")64 ap.add_argument("--limit", type=int, default=None)65 ap.add_argument("--text-meta-prefix", action="store_true", help="prepend <disc:..><era:..> tokens to the source text (M1 metadata control)")66 ap.add_argument("--overfit-test", action="store_true")67 ap.add_argument("--resume", action="store_true")68 ap.add_argument("--no-copy", action="store_true")69 ap.add_argument("--init-from", default=None, help="checkpoint dir (e.g. a denoising-pretrained run's best/) to initialise weights from; shapes must match")70 a = ap.parse_args()71 torch.manual_seed(a.seed)72 preset = PRESETS[a.preset]73 run_dir = ROOT / "artifacts/runs" / a.run74 run_dir.mkdir(parents=True, exist_ok=True)75 tok = HTRTokenizer(Path(a.tokenizer) / "tokenizer.model")76 splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}77 tr_docs = {d for d, s in splits.items() if s == "train"}78 va_docs = {d for d, s in splits.items() if s == "val"}79 variants = set(a.variants.split(",")) if a.variants else None80 all_train_pairs = [p for p in read_jsonl(a.pairs) if p["doc_id"] in tr_docs and (variants is None or p["variant"] in variants)]81 style_map = StyleLabelMap.from_pairs(all_train_pairs) if preset["arch"] == "m3" else None82 if style_map:83 style_map.save(run_dir / "style_classes.json")84 ds_tr = PairDataset(Path(a.pairs), tok, style_map=style_map, text_meta_prefix=a.text_meta_prefix, limit=a.limit, split_doc_ids=tr_docs, variants=variants)85 ds_va = PairDataset(Path(a.pairs), tok, style_map=style_map, text_meta_prefix=a.text_meta_prefix, limit=(a.limit // 5 if a.limit else None), split_doc_ids=va_docs, variants=variants)86 cfg = make_config(a.size, preset["arch"], tok.vocab_size, n_styles=(len(style_map.classes) if style_map else 0), use_copy=not a.no_copy)87 model = HTRModel(cfg)88 init_info = None89 if a.init_from:90 from safetensors.torch import load_file9192 state = load_file(str(Path(a.init_from) / "model.safetensors"))93 own = model.state_dict()94 loadable = {k: v for k, v in state.items() if k in own and own[k].shape == v.shape}95 missing = [k for k in own if k not in loadable and k != "out.weight"]96 model.load_state_dict({**own, **loadable})97 init_info = {"init_from": a.init_from, "loaded": len(loadable), "not_loaded": missing[:20], "n_not_loaded": len(missing)}98 print(f"initialised {len(loadable)} tensors from {a.init_from}; {len(missing)} left at random init (e.g. new style heads)")99 print(f"model {a.preset}/{a.size}: {model.n_params/1e6:.1f}M params; train pairs {len(ds_tr)} val pairs {len(ds_va)}; style classes {len(style_map.classes) if style_map else 0}")100 shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / "tokenizer.model")101 if a.overfit_test:102 batch = collate([ds_tr[i] for i in range(min(8, len(ds_tr)))], tok.pad_id)103 res = overfit_one_batch(model, batch, steps=150, device=a.device, loss_w=LossWeights(**preset["loss"], label_smoothing=0.0))104 write_json(run_dir / "overfit_test.json", {**run_header(), **res, "preset": a.preset, "size": a.size})105 print("overfit test:", res)106 return107 tcfg = TrainConfig(lr=a.lr, warmup_steps=a.warmup, max_steps=a.max_steps, accum=a.accum, eval_every=a.eval_every, patience=a.patience, amp=a.amp, seed=a.seed, max_tokens=a.max_tokens, max_batch=a.max_batch, loss=LossWeights(**preset["loss"]), device=a.device)108 tl = make_loader(ds_tr, a.max_tokens, a.max_batch, a.seed, True, num_workers=a.num_workers)109 vl = make_loader(ds_va, a.max_tokens, a.max_batch, 0, False, num_workers=a.num_workers)110 write_json(run_dir / "manifest.json", {**run_header({"run": a.run}), "args": vars(a), "preset": preset, "init": init_info, "model_config": cfg.to_dict(), "n_params": model.n_params, "n_train_pairs": len(ds_tr), "n_val_pairs": len(ds_va), "train_docs": len(tr_docs), "val_docs": len(va_docs), "style_classes": style_map.classes if style_map else None})111 trainer = Trainer(model, tcfg, run_dir, tl, vl)112 trainer.extra_files = [Path(a.tokenizer) / "tokenizer.model"]113 if a.resume and trainer.resume():114 print(f"resumed at step {trainer.step}")115 summary = trainer.train()116 for tag in ("best", "last"):117 if (run_dir / tag).exists():118 shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / tag / "tokenizer.model")119 write_json(run_dir / tag / "release_meta.json", {"name": f"htr-{a.preset}-{a.size}", "text_meta_prefix": a.text_meta_prefix, "style_classes": style_map.classes if style_map else None, "run": a.run})120 print("summary:", summary)121122123if __name__ == "__main__":124 main()125