spb/wpai0002_uqo

1#!/usr/bin/env python32"""Minimum-risk fine-tuning of an MLE-fine-tuned HTR checkpoint (see htr.training.mrt). Model selection by greedy chrF++3on a fixed validation slice (not CE). Writes artifacts/runs/<run>/{best,last,train_log.jsonl,train_summary.json}."""4from __future__ import annotations56import argparse7import math8import shutil9import sys10import time11from pathlib import Path1213import numpy as np14import torch1516ROOT = Path(__file__).resolve().parents[1]17sys.path.insert(0, str(ROOT / "src"))18from htr.evaluation.metrics import _chrf  # noqa: E40219from htr.models.losses import LossWeights  # noqa: E40220from htr.models.tokenizer import HTRTokenizer  # noqa: E40221from htr.training.data import PairDataset, StyleLabelMap, make_loader  # noqa: E40222from htr.training.mrt import Discriminator, mrt_loss  # noqa: E40223from htr.training.trainer import _to, load_checkpoint, save_checkpoint  # noqa: E40224from htr.utils.io import append_jsonl, read_jsonl, write_json  # noqa: E40225from htr.utils.provenance import run_header  # noqa: E402262728def main() -> None:29    ap = argparse.ArgumentParser()30    ap.add_argument("--pairs", required=True)31    ap.add_argument("--splits", required=True)32    ap.add_argument("--init-from", required=True, help="MLE-fine-tuned checkpoint dir (…/best)")33    ap.add_argument("--run", required=True)34    ap.add_argument("--max-steps", type=int, default=2000)35    ap.add_argument("--eval-every", type=int, default=200)36    ap.add_argument("--n-val", type=int, default=600)37    ap.add_argument("--lr", type=float, default=1e-5)38    ap.add_argument("--warmup", type=int, default=100)39    ap.add_argument("--max-tokens", type=int, default=3072)40    ap.add_argument("--max-batch", type=int, default=16)41    ap.add_argument("--n-samples", type=int, default=8)42    ap.add_argument("--temperature", type=float, default=0.8)43    ap.add_argument("--top-p", type=float, default=0.95)44    ap.add_argument("--alpha", type=float, default=0.005, help="Q(y) ∝ p(y|x)^alpha; Shen et al. 2016 use 5e-3 — with 0.5 the distribution is one-hot on greedy (no gradient signal)")45    ap.add_argument("--w-ce", type=float, default=0.3)46    ap.add_argument("--fact-weight", type=float, default=0.5)47    ap.add_argument("--no-gold", action="store_true")48    ap.add_argument("--disc", default=None, help="discriminator dir (…/best) → reward term disc_weight·(1 − P(original))")49    ap.add_argument("--disc-weight", type=float, default=0.0)50    ap.add_argument("--patience", type=int, default=5)51    ap.add_argument("--seed", type=int, default=1)52    ap.add_argument("--device", default="auto")53    ap.add_argument("--amp", default="auto")54    ap.add_argument("--num-workers", type=int, default=0)55    a = ap.parse_args()56    torch.manual_seed(a.seed)57    dev = torch.device(a.device if a.device != "auto" else ("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")))58    if dev.type == "cuda" and hasattr(torch.backends.cuda, "enable_cudnn_sdp"):59        torch.backends.cuda.enable_cudnn_sdp(False)60    init = Path(a.init_from)61    model = load_checkpoint(init, dev)62    tok = HTRTokenizer(init / "tokenizer.model")63    run_dir = ROOT / "artifacts/runs" / a.run64    run_dir.mkdir(parents=True, exist_ok=True)65    splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}66    tr_ids = {d for d, s in splits.items() if s == "train"}67    va_ids = {d for d, s in splits.items() if s == "val"}68    style_map = StyleLabelMap.load(init / "style_map.json") if (init / "style_map.json").exists() else None69    ds_tr = PairDataset(Path(a.pairs), tok, style_map=style_map, split_doc_ids=tr_ids)70    tl = make_loader(ds_tr, a.max_tokens, a.max_batch, a.seed, True, num_workers=a.num_workers)71    val_pairs = [p for p in read_jsonl(a.pairs) if p["doc_id"] in va_ids][: a.n_val]72    from htr.inference.pipeline import HTRReconstructor7374    disc = Discriminator(a.disc, dev) if a.disc and a.disc_weight > 0 else None75    use_bf16 = (a.amp == "bf16") or (a.amp == "auto" and dev.type == "cuda")76    ac = torch.autocast(device_type=dev.type, dtype=torch.bfloat16) if use_bf16 else None7778    def val_chrf() -> dict:79        rec = HTRReconstructor(model, tok, {"name": a.run})80        outs = rec.reconstruct_many([p["source_text"] for p in val_pairs], [p["metadata"] for p in val_pairs], batch_size=48)81        model.train()82        sc = [float(_chrf.sentence_score(o.reconstruction, [p["target_text"]]).score) for o, p in zip(outs, val_pairs)]83        idn = [float(_chrf.sentence_score(p["source_text"], [p["target_text"]]).score) for p in val_pairs]84        nr = float(np.mean([o.preservation["number_recall"] for o in outs]))85        ev = {"val_chrf": float(np.mean(sc)), "val_delta_identity": float(np.mean(sc) - np.mean(idn)), "val_number_recall": nr}86        if disc is not None:87            ev["val_disc_p_human"] = float(np.mean(disc.p_human([o.reconstruction for o in outs])))88        return ev8990    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01, betas=(0.9, 0.98))91    ev0 = val_chrf()92    print("init", ev0, flush=True)93    append_jsonl(run_dir / "train_log.jsonl", {"step": 0, "eval": ev0})94    best, bad, step, t0 = ev0["val_chrf"], 0, 0, time.time()95    save_checkpoint(model, run_dir / "best", {"step": 0, "mrt": True, **ev0})96    shutil.copy(init / "tokenizer.model", run_dir / "best" / "tokenizer.model")97    for f in ("release_meta.json", "style_map.json"):98        if (init / f).exists():99            shutil.copy(init / f, run_dir / "best" / f)100    write_json(run_dir / "manifest.json", {**run_header({"run": a.run, "stage": "mrt"}), "args": vars(a), "init_from": str(init)})101    model.train()102    done = False103    while not done:104        for batch in tl:105            b = _to(batch, dev)106            for g in opt.param_groups:107                g["lr"] = a.lr * min(1.0, (step + 1) / a.warmup) * 0.5 * (1 + math.cos(math.pi * min(1.0, step / a.max_steps)))108            with (ac or torch.autocast(device_type="cpu", enabled=False)):109                loss, logs = mrt_loss(model, tok, b, a.n_samples, a.temperature, a.top_p, a.alpha, a.w_ce, a.fact_weight, not a.no_gold, LossWeights(fact=0.5), disc=disc, disc_weight=a.disc_weight)110            loss.backward()111            gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)112            opt.step()113            opt.zero_grad(set_to_none=True)114            step += 1115            if step % 10 == 0:116                append_jsonl(run_dir / "train_log.jsonl", {"step": step, "grad_norm": float(gn), "elapsed_s": round(time.time() - t0, 1), **logs})117            if step % a.eval_every == 0 or step >= a.max_steps:118                ev = val_chrf()119                append_jsonl(run_dir / "train_log.jsonl", {"step": step, "eval": ev})120                print(step, ev, flush=True)121                if ev["val_chrf"] > best:122                    best, bad = ev["val_chrf"], 0123                    save_checkpoint(model, run_dir / "best", {"step": step, "mrt": True, **ev})124                    shutil.copy(init / "tokenizer.model", run_dir / "best" / "tokenizer.model")125                    for f in ("release_meta.json", "style_map.json"):126                        if (init / f).exists():127                            shutil.copy(init / f, run_dir / "best" / f)128                else:129                    bad += 1130            if step >= a.max_steps or bad >= a.patience:131                done = True132                break133    write_json(run_dir / "train_summary.json", {**run_header({"run": a.run, "stage": "mrt"}), "steps": step, "best_val_chrf": best, "init_val_chrf": ev0["val_chrf"], "wall_seconds": round(time.time() - t0, 1), "args": vars(a)})134    print("done: best val chrF", best, "(init", ev0["val_chrf"], ")")135136137if __name__ == "__main__":138    main()139