spb/wpai0002_uqo

1#!/usr/bin/env python32"""Denoising pretraining on train-split original passages (no API spend). Produces a checkpoint usable with3``scripts/train.py --init-from <run>/best`` for pair fine-tuning."""4from __future__ import annotations56import argparse7import sys8from pathlib import Path910ROOT = Path(__file__).resolve().parents[1]11sys.path.insert(0, str(ROOT / "src"))12import torch  # noqa: E4021314from htr.models.config import make_config  # noqa: E40215from htr.models.losses import LossWeights  # noqa: E40216from htr.models.tokenizer import HTRTokenizer  # noqa: E40217from htr.models.transformer import HTRModel  # noqa: E40218from htr.training.data import BucketSampler, collate  # noqa: E40219from htr.training.denoise import DenoiseDataset  # noqa: E40220from htr.training.trainer import TrainConfig, Trainer  # noqa: E40221from htr.utils.io import read_jsonl, write_json  # noqa: E40222from htr.utils.provenance import run_header  # noqa: E402232425def main() -> None:26    ap = argparse.ArgumentParser()27    ap.add_argument("--corpus", required=True, help="parsed corpus dir (documents.jsonl + passages.jsonl)")28    ap.add_argument("--splits", required=True)29    ap.add_argument("--tokenizer", required=True)30    ap.add_argument("--run", required=True)31    ap.add_argument("--arch", default="m1", choices=["m1", "m2", "m3"])32    ap.add_argument("--size", default="tiny")33    ap.add_argument("--max-steps", type=int, default=20000)34    ap.add_argument("--eval-every", type=int, default=500)35    ap.add_argument("--lr", type=float, default=5e-4)36    ap.add_argument("--warmup", type=int, default=500)37    ap.add_argument("--max-tokens", type=int, default=4096)38    ap.add_argument("--max-batch", type=int, default=32)39    ap.add_argument("--patience", type=int, default=8)40    ap.add_argument("--seed", type=int, default=1)41    ap.add_argument("--device", default="auto")42    ap.add_argument("--limit", type=int, default=None)43    ap.add_argument("--variants-per-passage", type=int, default=1)44    ap.add_argument("--n-styles", type=int, default=8)45    ap.add_argument("--resume", action="store_true")46    ap.add_argument("--init-from", default=None, help="initialise weights from a checkpoint dir (fresh optimizer/schedule), e.g. to continue after a divergence with a lower lr")47    ap.add_argument("--num-workers", type=int, default=0, help="DataLoader workers (noising + tokenisation are CPU-bound; use ≥ 8 on a GPU box)")48    ap.add_argument("--amp", default="auto", help="auto | bf16 | none (auto = bf16 on CUDA, fp32 on MPS/CPU)")49    a = ap.parse_args()50    torch.manual_seed(a.seed)51    run_dir = ROOT / "artifacts/runs" / a.run52    run_dir.mkdir(parents=True, exist_ok=True)53    tok = HTRTokenizer(Path(a.tokenizer) / "tokenizer.model")54    splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}55    tr = {d for d, s in splits.items() if s == "train"}56    va = {d for d, s in splits.items() if s == "val"}57    corpus = Path(a.corpus)58    ds_tr = DenoiseDataset(str(corpus / "passages.jsonl"), str(corpus / "documents.jsonl"), tok, tr, seed=a.seed, limit=a.limit, variants_per_passage=a.variants_per_passage)59    ds_va = DenoiseDataset(str(corpus / "passages.jsonl"), str(corpus / "documents.jsonl"), tok, va, seed=a.seed + 1, limit=(a.limit // 10 if a.limit else 1500))60    cfg = make_config(a.size, a.arch, tok.vocab_size, n_styles=a.n_styles)61    model = HTRModel(cfg)62    if a.init_from:63        from htr.training.trainer import load_checkpoint6465        init = load_checkpoint(Path(a.init_from), "cpu")66        missing, unexpected = model.load_state_dict(init.state_dict(), strict=False)67        print(f"init from {a.init_from}: missing {len(missing)} unexpected {len(unexpected)}")68    print(f"denoise {a.arch}/{a.size}: {model.n_params/1e6:.1f}M params; train items {len(ds_tr)} val items {len(ds_va)}")69    lengths_tr = [min(512, len(it["text"]) // 4 + 8) for it in ds_tr.items]70    lengths_va = [min(512, len(it["text"]) // 4 + 8) for it in ds_va.items]71    import functools7273    coll = functools.partial(collate, pad_id=tok.pad_id)74    pf = 4 if a.num_workers else None75    tl = torch.utils.data.DataLoader(ds_tr, batch_sampler=BucketSampler(lengths_tr, a.max_tokens, a.max_batch, a.seed, True), collate_fn=coll, num_workers=a.num_workers, prefetch_factor=pf)76    vl = torch.utils.data.DataLoader(ds_va, batch_sampler=BucketSampler(lengths_va, a.max_tokens, a.max_batch, 0, False), collate_fn=coll, num_workers=a.num_workers, prefetch_factor=pf)77    tcfg = TrainConfig(lr=a.lr, warmup_steps=a.warmup, max_steps=a.max_steps, 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(fact=0.5), device=a.device)78    import shutil7980    shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / "tokenizer.model")81    write_json(run_dir / "manifest.json", {**run_header({"run": a.run, "stage": "denoise_pretrain"}), "args": vars(a), "model_config": cfg.to_dict(), "n_params": model.n_params, "init_from": a.init_from, "n_train_items": len(ds_tr), "n_val_items": len(ds_va)})82    trainer = Trainer(model, tcfg, run_dir, tl, vl)83    trainer.extra_files = [Path(a.tokenizer) / "tokenizer.model"]84    if a.resume and trainer.resume():85        print(f"resumed at {trainer.step}")86    s = trainer.train()87    for tag in ("best", "last"):88        if (run_dir / tag).exists():89            shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / tag / "tokenizer.model")90            write_json(run_dir / tag / "release_meta.json", {"name": f"htr-denoise-{a.arch}-{a.size}", "stage": "denoise_pretrain", "run": a.run})91    print("summary:", s)929394if __name__ == "__main__":95    main()96