spb/wpai0002_uqo

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}394041def main() -> None:42    ap = argparse.ArgumentParser()43    ap.add_argument("--pairs", required=True)44    ap.add_argument("--splits", required=True)45    ap.add_argument("--tokenizer", required=True, help="dir with tokenizer.model")46    ap.add_argument("--run", required=True, help="run name under artifacts/runs/")47    ap.add_argument("--preset", default="m1", choices=sorted(PRESETS))48    ap.add_argument("--size", default="tiny")49    ap.add_argument("--max-steps", type=int, default=5000)50    ap.add_argument("--eval-every", type=int, default=250)51    ap.add_argument("--lr", type=float, default=3e-4)52    ap.add_argument("--warmup", type=int, default=300)53    ap.add_argument("--max-tokens", type=int, default=8192)54    ap.add_argument("--max-batch", type=int, default=48)55    ap.add_argument("--accum", type=int, default=1)56    ap.add_argument("--patience", type=int, default=6)57    ap.add_argument("--seed", type=int, default=1)58    ap.add_argument("--device", default="auto")59    ap.add_argument("--amp", default="auto")60    ap.add_argument("--variants", default=None, help="restrict training variants, e.g. A,B,C,D (hold out E)")61    ap.add_argument("--limit", type=int, default=None)62    ap.add_argument("--text-meta-prefix", action="store_true", help="prepend <disc:..><era:..> tokens to the source text (M1 metadata control)")63    ap.add_argument("--overfit-test", action="store_true")64    ap.add_argument("--resume", action="store_true")65    ap.add_argument("--no-copy", action="store_true")66    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")67    a = ap.parse_args()68    torch.manual_seed(a.seed)69    preset = PRESETS[a.preset]70    run_dir = ROOT / "artifacts/runs" / a.run71    run_dir.mkdir(parents=True, exist_ok=True)72    tok = HTRTokenizer(Path(a.tokenizer) / "tokenizer.model")73    splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}74    tr_docs = {d for d, s in splits.items() if s == "train"}75    va_docs = {d for d, s in splits.items() if s == "val"}76    variants = set(a.variants.split(",")) if a.variants else None77    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)]78    style_map = StyleLabelMap.from_pairs(all_train_pairs) if preset["arch"] == "m3" else None79    if style_map:80        style_map.save(run_dir / "style_classes.json")81    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)82    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)83    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)84    model = HTRModel(cfg)85    init_info = None86    if a.init_from:87        from safetensors.torch import load_file8889        state = load_file(str(Path(a.init_from) / "model.safetensors"))90        own = model.state_dict()91        loadable = {k: v for k, v in state.items() if k in own and own[k].shape == v.shape}92        missing = [k for k in own if k not in loadable and k != "out.weight"]93        model.load_state_dict({**own, **loadable})94        init_info = {"init_from": a.init_from, "loaded": len(loadable), "not_loaded": missing[:20], "n_not_loaded": len(missing)}95        print(f"initialised {len(loadable)} tensors from {a.init_from}; {len(missing)} left at random init (e.g. new style heads)")96    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}")97    shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / "tokenizer.model")98    if a.overfit_test:99        batch = collate([ds_tr[i] for i in range(min(8, len(ds_tr)))], tok.pad_id)100        res = overfit_one_batch(model, batch, steps=150, device=a.device, loss_w=LossWeights(**preset["loss"], label_smoothing=0.0))101        write_json(run_dir / "overfit_test.json", {**run_header(), **res, "preset": a.preset, "size": a.size})102        print("overfit test:", res)103        return104    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)105    tl = make_loader(ds_tr, a.max_tokens, a.max_batch, a.seed, True)106    vl = make_loader(ds_va, a.max_tokens, a.max_batch, 0, False)107    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})108    trainer = Trainer(model, tcfg, run_dir, tl, vl)109    if a.resume and trainer.resume():110        print(f"resumed at step {trainer.step}")111    summary = trainer.train()112    for tag in ("best", "last"):113        if (run_dir / tag).exists():114            shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / tag / "tokenizer.model")115            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})116    print("summary:", summary)117118119if __name__ == "__main__":120    main()121