spb/wpai0002_uqo

1"""Package a trained checkpoint into a self-contained, independently loadable release directory:23    <release>/model.safetensors, config.json, tokenizer.model, tokenizer_config.json, release_meta.json,4             generation_config.json, README.md (model card), LICENSE-MODEL.txt, MANIFEST.json (sha256 of every file)56The release is loadable with ``htr.inference.pipeline.HTRReconstructor.from_pretrained(<release>)`` and by the7``htr`` CLI; it carries no training data. A Hugging Face-compatible layout (config.json + model.safetensors +8tokenizer files + README.md with YAML front-matter) is used so that the directory can be uploaded as-is once the9author approves publication (never automatic)."""10from __future__ import annotations1112import json13import shutil14from pathlib import Path1516from htr.utils.hashing import sha256_bytes17from htr.utils.io import read_json, write_json18from htr.utils.provenance import run_header, utc_now1920ROOT = Path(__file__).resolve().parents[3]212223def _sha(path: Path) -> str:24    return sha256_bytes(path.read_bytes())252627def export_release(checkpoint_dir: Path, out_dir: Path, name: str, version: str, metrics: dict | None = None, extra_meta: dict | None = None, card_template: Path | None = None) -> dict:28    checkpoint_dir, out_dir = Path(checkpoint_dir), Path(out_dir)29    out_dir.mkdir(parents=True, exist_ok=True)30    for f in ("model.safetensors", "config.json", "tokenizer.model"):31        shutil.copy(checkpoint_dir / f, out_dir / f)32    tok_cfg = checkpoint_dir.parent.parent.parent / "tokenizers"  # best effort; tokenizer_config lives next to the trained tokenizer33    cfg = read_json(checkpoint_dir / "config.json")34    rel_meta = read_json(checkpoint_dir / "release_meta.json") if (checkpoint_dir / "release_meta.json").exists() else {}35    ck_meta = read_json(checkpoint_dir / "checkpoint_meta.json") if (checkpoint_dir / "checkpoint_meta.json").exists() else {}36    meta = {**rel_meta, "name": name, "version": version, "exported_at": utc_now(), "arch": cfg["arch"], "n_params": cfg.get("n_params"), "source_checkpoint": str(checkpoint_dir), "checkpoint_meta": ck_meta, "metrics": metrics or {}, **(extra_meta or {})}37    write_json(out_dir / "release_meta.json", meta)38    write_json(out_dir / "tokenizer_config.json", {"type": "sentencepiece", "file": "tokenizer.model", "pad_id": 0, "bos_id": 1, "eos_id": 2, "unk_id": 3})39    write_json(out_dir / "generation_config.json", {"decoding": "greedy", "beam": 1, "max_new_tokens": "1.6 x input tokens + 16 (capped by config.max_len)", "deterministic": True, "placeholder_protocol": "citations/math/urls are replaced by ⟦Cn⟧/⟦Mn⟧/⟦Un⟧ before encoding and restored after decoding"})40    card = render_model_card(meta, cfg, card_template)41    (out_dir / "README.md").write_text(card)42    (out_dir / "LICENSE-MODEL.txt").write_text(MODEL_LICENSE_NOTE)43    manifest = {**run_header({"release": name, "version": version}), "files": {f.name: {"sha256": _sha(f), "bytes": f.stat().st_size} for f in sorted(out_dir.iterdir()) if f.is_file() and f.name != "MANIFEST.json"}}44    write_json(out_dir / "MANIFEST.json", manifest)45    return manifest464748MODEL_LICENSE_NOTE = """HTR model weights and code: Apache License 2.0 (see repository LICENSE).49Training data: documents under CC0 / public domain / CC BY / CC BY-SA / CC BY-NC / CC BY-NC-SA licences (see the data card and50docs/data_governance/LICENSE_MATRIX.md). Because CC BY-NC material may be among the training sources, the weights are released51for NON-COMMERCIAL RESEARCH USE unless the release notes state that the checkpoint was trained on the CC0/CC BY/public-domain subset only.52Outputs are synthetic reconstructions; they must not be presented as the original text or as evidence of human authorship.53"""545556def render_model_card(meta: dict, cfg: dict, template: Path | None = None) -> str:57    tpl_path = template or (ROOT / "model_cards" / "TEMPLATE.md")58    tpl = tpl_path.read_text() if tpl_path.exists() else "# {name}\n"59    m = meta.get("metrics", {})60    fields = {61        "name": meta["name"], "version": meta["version"], "arch": cfg["arch"], "n_params_m": f"{(cfg.get('n_params') or 0)/1e6:.1f}",62        "d_model": cfg["d_model"], "layers": f"{cfg['n_enc_layers']}+{cfg['n_dec_layers']}", "vocab": cfg["vocab_size"], "max_len": cfg["max_len"],63        "exported_at": meta["exported_at"], "run": meta.get("run", "?"), "metrics_json": json.dumps(m, indent=2) if m else "pending",64        "style_classes": len(meta.get("style_classes") or []) or "n/a", "training_data": meta.get("training_data", "see release notes"),65    }66    out = tpl67    for k, v in fields.items():68        out = out.replace("{" + k + "}", str(v))69    return out70