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    a = ap.parse_args()47    torch.manual_seed(a.seed)48    run_dir = ROOT / "artifacts/runs" / a.run49    run_dir.mkdir(parents=True, exist_ok=True)50    tok = HTRTokenizer(Path(a.tokenizer) / "tokenizer.model")51    splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}52    tr = {d for d, s in splits.items() if s == "train"}53    va = {d for d, s in splits.items() if s == "val"}54    corpus = Path(a.corpus)55    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)56    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))57    cfg = make_config(a.size, a.arch, tok.vocab_size, n_styles=a.n_styles)58    model = HTRModel(cfg)59    print(f"denoise {a.arch}/{a.size}: {model.n_params/1e6:.1f}M params; train items {len(ds_tr)} val items {len(ds_va)}")60    lengths_tr = [min(512, len(it["text"]) // 4 + 8) for it in ds_tr.items]61    lengths_va = [min(512, len(it["text"]) // 4 + 8) for it in ds_va.items]62    tl = torch.utils.data.DataLoader(ds_tr, batch_sampler=BucketSampler(lengths_tr, a.max_tokens, a.max_batch, a.seed, True), collate_fn=lambda b: collate(b, tok.pad_id))63    vl = torch.utils.data.DataLoader(ds_va, batch_sampler=BucketSampler(lengths_va, a.max_tokens, a.max_batch, 0, False), collate_fn=lambda b: collate(b, tok.pad_id))64    tcfg = TrainConfig(lr=a.lr, warmup_steps=a.warmup, max_steps=a.max_steps, eval_every=a.eval_every, patience=a.patience, seed=a.seed, max_tokens=a.max_tokens, max_batch=a.max_batch, loss=LossWeights(fact=0.5), device=a.device)65    import shutil6667    shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / "tokenizer.model")68    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, "n_train_items": len(ds_tr), "n_val_items": len(ds_va)})69    trainer = Trainer(model, tcfg, run_dir, tl, vl)70    if a.resume and trainer.resume():71        print(f"resumed at {trainer.step}")72    s = trainer.train()73    for tag in ("best", "last"):74        if (run_dir / tag).exists():75            shutil.copy(Path(a.tokenizer) / "tokenizer.model", run_dir / tag / "tokenizer.model")76            write_json(run_dir / tag / "release_meta.json", {"name": f"htr-denoise-{a.arch}-{a.size}", "stage": "denoise_pretrain", "run": a.run})77    print("summary:", s)787980if __name__ == "__main__":81    main()82