code/export_hf.py

"""Package a trained run for the private HF repo and (only on request) push it (classifier-plan §3.5).

Run on this host: ``python -m training.export_hf RUN_DIR --out DIR [--metrics F ...] [--push --yes]``.

The package holds: ``adapter/`` (LoRA; ``base_model_name_or_path`` rewritten to the hub id), ``heads.safetensors``,
``taxonomy.json``, ``temperatures.json``, ``config.json`` (training settings without local paths), the metrics and
``REPORT.md`` files given, a generated model card ``README.md`` and ``code/`` (the training package and its
requirements). Never data, predictions or the ledger: the package is checked for ``*.jsonl`` and ``ledger*`` files
before anything is written or pushed. The push needs both ``--push`` and ``--yes``, reads ``HF_TOKEN`` from the
environment (never printed) and creates the repo private; the HF commit becomes ``model_revision`` (§6.3).
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any, Callable

DEFAULT_REPO = "echoforest/naver-psy-classifier-gemma4-12b"
BASES = {"base": "google/gemma-4-12B", "it": "google/gemma-4-12B-it"}
PATH_KEYS = ("data", "model", "out", "resume_from")
FORBIDDEN_SUFFIXES = (".jsonl", ".jsonl.tmp", ".pt", ".tar", ".age")
HERE = Path(__file__).resolve().parent


def base_model_id(cfg: dict) -> str:
    if cfg.get("base_model"):
        return cfg["base_model"]
    return BASES[cfg.get("checkpoint", "base")]


def clean_config(cfg: dict) -> dict:
    """Training settings without local paths (keys named in PATH_KEYS or absolute-path string values)."""
    out = {}
    for k, v in cfg.items():
        if k in PATH_KEYS or (isinstance(v, str) and (v.startswith("/") or v.startswith("~"))):
            continue
        out[k] = v
    out["base_model"] = base_model_id(cfg)
    return out


def forbidden(folder: Path) -> list[str]:
    """Files that must never be in a package: data, predictions, optimizer state, the ledger."""
    bad = []
    for p in folder.rglob("*"):
        name = p.name.lower()
        ledger = name.startswith("ledger") and not name.endswith(".py")  # the ledger file, not code/ledger.py
        if p.is_file() and (name.endswith(FORBIDDEN_SUFFIXES) or ledger or "predictions" in p.relative_to(folder).parts):
            bad.append(str(p.relative_to(folder)))
    return sorted(bad)


def model_card(cfg: dict, metrics: dict[str, Any], repo: str) -> str:
    lines = [
        "---", "license: apache-2.0", f"base_model: {base_model_id(cfg)}", "library_name: peft", "private: true", "---",
        "", f"# {repo.split('/')[-1]}", "",
        "LoRA adapter and per-field linear heads on the Gemma 4 12B text decoder for the "
        f"`{cfg.get('taxonomy_version') or 'nsb-v1'}` taxonomy of",
        "Naver 종목토론실 posts (naver-psy, docs/classifier-plan.md). Private; not for redistribution.", "",
        "- Agreement figures are agreement with the teacher labels, not accuracy against ground truth.",
        "- Input: `nsb-input-v1` render, segments joined with newlines, suffix `### LABELS`, one BOS.",
        "- Output: raw logits per field; calibrated with `temperatures.json`, then gated (training/taxonomy.py).",
        "- No training data, predictions or post text are in this repository.", "",
        "| Setting | Value |", "|---|---|",
    ]
    for k in ("arm", "checkpoint", "pool", "epochs", "max_len", "lora_r", "lora_alpha", "lr_lora", "lr_heads", "seed",
              "data_version", "data_manifest_sha256"):
        if k in cfg:
            lines.append(f"| {k} | `{cfg[k]}` |")
    if metrics:
        lines += ["", "## Metrics", ""]
        for name, value in metrics.items():
            lines += [f"### {name}", "", "```json", json.dumps(value, indent=1, ensure_ascii=False)[:4000], "```", ""]
    return "\n".join(lines) + "\n"


def copy_code(dst: Path) -> None:
    dst.mkdir(parents=True)
    for p in sorted(HERE.glob("*.py")) + sorted(HERE.glob("requirements*.txt")):
        shutil.copy(p, dst / p.name)
    runpod = HERE / "runpod"
    if runpod.is_dir():
        (dst / "runpod").mkdir()
        for p in sorted(runpod.iterdir()):
            if p.is_file() and not p.name.endswith(".pyc"):
                shutil.copy(p, dst / "runpod" / p.name)


def package(run: Path, out: Path, metrics: list[Path] = (), which: str = "best", repo: str = DEFAULT_REPO) -> Path:
    run, out = Path(run), Path(out)
    if out.exists():
        raise SystemExit(f"refusing to overwrite {out}")
    src = run / which
    cfg = json.loads((run / "config.json").read_text())
    for need in (src / "adapter" / "adapter_config.json", src / "adapter" / "adapter_model.safetensors",
                 src / "heads.safetensors", run / "taxonomy.json", run / "temperatures.json"):
        if not need.exists():
            raise SystemExit(f"missing {need}")
    tmp = out.with_name(out.name + ".tmp")
    shutil.rmtree(tmp, ignore_errors=True)
    (tmp / "adapter").mkdir(parents=True)
    shutil.copy(src / "adapter" / "adapter_model.safetensors", tmp / "adapter" / "adapter_model.safetensors")
    adapter_cfg = json.loads((src / "adapter" / "adapter_config.json").read_text())
    adapter_cfg["base_model_name_or_path"] = base_model_id(cfg)
    (tmp / "adapter" / "adapter_config.json").write_text(json.dumps(adapter_cfg, indent=2) + "\n")
    shutil.copy(src / "heads.safetensors", tmp / "heads.safetensors")
    shutil.copy(run / "taxonomy.json", tmp / "taxonomy.json")
    shutil.copy(run / "temperatures.json", tmp / "temperatures.json")
    settings = clean_config(cfg)
    state = src / "state.json"
    if state.exists():
        settings["checkpoint_step"] = json.loads(state.read_text()).get("step")
    (tmp / "config.json").write_text(json.dumps(settings, indent=1, ensure_ascii=False) + "\n")
    loaded = {}
    for m in [Path(x) for x in metrics]:
        if m.suffix == ".json":
            loaded[m.stem] = json.loads(m.read_text())
        shutil.copy(m, tmp / m.name)
    if (run / "REPORT.md").exists() and not (tmp / "REPORT.md").exists():
        shutil.copy(run / "REPORT.md", tmp / "REPORT.md")
    (tmp / "README.md").write_text(model_card(settings, loaded, repo))
    copy_code(tmp / "code")
    bad = forbidden(tmp)
    if bad:
        shutil.rmtree(tmp)
        raise SystemExit(f"refusing to package data/predictions/ledger files: {bad[:5]}")
    tmp.rename(out)
    return out


def push(folder: Path, repo: str = DEFAULT_REPO, api_factory: Callable[..., Any] | None = None) -> str:
    """Upload to the private repo; returns ``repo@<commit>`` (the ``model_revision``)."""
    folder = Path(folder)
    bad = forbidden(folder)
    if bad:
        raise SystemExit(f"refusing to push data/predictions/ledger files: {bad[:5]}")
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise SystemExit("HF_TOKEN is not set")
    if api_factory is None:
        from huggingface_hub import HfApi as api_factory
    api = api_factory(token=token)
    api.create_repo(repo, private=True, exist_ok=True, repo_type="model")
    info = api.upload_folder(folder_path=str(folder), repo_id=repo, repo_type="model",
                             commit_message="Upload naver-psy classifier (adapter, heads, calibration)")
    oid = getattr(info, "oid", None) or getattr(info, "commit_id", None)
    if not oid or len(oid) != 40:
        raise SystemExit("upload returned no 40-hex commit id")
    return f"{repo}@{oid}"


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("run")
    p.add_argument("--out", required=True)
    p.add_argument("--which", default="best", choices=["best", "last"])
    p.add_argument("--metrics", nargs="*", default=[])
    p.add_argument("--repo", default=DEFAULT_REPO)
    p.add_argument("--push", action="store_true", help="Upload to the private HF repo (also needs --yes)")
    p.add_argument("--yes", action="store_true")
    args = p.parse_args(argv)
    if args.push and not args.yes:
        p.error("--push also needs --yes")
    out = package(Path(args.run), Path(args.out), [Path(m) for m in args.metrics], args.which, args.repo)
    print(json.dumps({"packaged": str(out), "files": sorted(str(x.relative_to(out)) for x in out.rglob("*")
                                                            if x.is_file() and "code" not in x.parts)}))
    if args.push:
        revision = push(out, args.repo)
        print(json.dumps({"model_revision": revision}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())