Blame for source/scripts/deployment/prepare_mac_zh_v1.py

"""Package the committed main branch, vimeml-tiny-zh-v1 inference weights, PyTorch FP32 fixtures and the
development split of ime-standard-zh-v1-r1 with v1's reference scores, for Core ML and quantization on a Mac.

Blind items never enter the package: every benchmark file is filtered to split == development.
Brief: docs/reference/mac-zh-v1-coreml.md.
"""
import argparse
import json
import shutil
import subprocess
import sys
import tempfile
import zipfile
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))

CHECKPOINT = "artifacts/models/tiny-zh13-12k-fresh-s1-d384-l9/final.pt"
CHECKPOINT_SHA256 = "44c423a1e7cb0e32161afc94d6167cabb439a3d25fd9de6d1fbb150a927542d4"
DEPLOYMENT = "artifacts/models/tiny-zh13-12k-fresh-s1-d384-l9/deployment.pt"
TOKENIZER_DIR = "artifacts/tokenizers/zh-v13/zh-unigram-12288"
TOKENIZER_SHA256 = "68bd37c45bcb018b32b2eef5b241561ab5b6dc1b77c1c3aadcc1b41c4bf47e5a"
REFERENCE = "outputs/deployment/vimeml-tiny-zh-v1-fp32-input"
FROZEN = "outputs/benchmarks/ime-standard-zh-v1/frozen-r1"
FILTERED = ("items.jsonl", "pool-full.jsonl", "pool-abbrev.jsonl", "rime-full.jsonl", "rime-abbrev.jsonl",
            "cand-full.jsonl", "cand-abbrev.jsonl", "overlap/overlap.jsonl")
SCORES = "outputs/benchmarks/zh13-20m/runs/12k-fresh-s1"
RESULTS = ("outputs/benchmarks/zh13-20m/judgement/scores/12k-fresh-s1-full",
           "outputs/benchmarks/zh13-20m/judgement/scores/12k-fresh-s1-abbrev",
           "outputs/benchmarks/zh13-20m/judgement/scores/rime-full",
           "outputs/benchmarks/zh13-20m/judgement/scores/rime-abbrev")
EXPECTED_TOP1 = {"full": 1363, "abbrev": 337}
COPIED = (f"{TOKENIZER_DIR}/tokenizer.model", f"{TOKENIZER_DIR}/tokenizer.vocab",
          "artifacts/models/tiny-zh13-12k-fresh-s1-d384-l9/config.json",
          "artifacts/models/tiny-zh13-12k-fresh-s1-d384-l9/summary.json",
          "artifacts/models/tiny-zh13-12k-fresh-s1-d384-l9/environment.json")
TEXT = "你好，我想找一家离地铁站近一点的酒店，晚上想吃火锅，你有推荐的餐馆吗？明天上午我们去故宫参观，下午再去逛街。"


def git(*args):
    return subprocess.check_output(["git", *args], cwd=ROOT, text=True).strip()


def write_json(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", newline="\n")


def read_jsonl(path):
    with open(path, encoding="utf-8") as stream:
        return [json.loads(line) for line in stream if line.strip()]


def write_jsonl(path, rows):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8", newline="\n") as stream:
        for row in rows:
            stream.write(json.dumps(row, ensure_ascii=False) + "\n")


def prepare_weights(payload):
    import numpy as np
    import sentencepiece as spm
    import torch
    from vimeml.benchmarks.zh13_lm_scores import load
    from vimeml.training.data import file_sha

    torch.set_num_threads(4)
    if file_sha(ROOT / CHECKPOINT) != CHECKPOINT_SHA256 or file_sha(ROOT / TOKENIZER_DIR / "tokenizer.model") != TOKENIZER_SHA256:
        raise ValueError("Checkpoint or tokenizer differs from the confirmed vimeml-tiny-zh-v1")
    saved = torch.load(ROOT / CHECKPOINT, map_location="cpu", weights_only=True)
    if any(t.dtype != torch.float32 for t in saved["model"].values()):
        raise ValueError("Expected FP32 master weights")
    exported = {key: saved[key] for key in ("format", "architecture", "model", "model_config", "step", "epoch",
                                            "signatures", "tokenizer_identity")}
    exported["config"] = {key: saved["config"][key] for key in ("arm", "architecture", "stage")}
    exported.update(deployment_only=True, source_checkpoint_sha256=CHECKPOINT_SHA256, release_name="vimeml-tiny-zh-v1")
    destination = payload / DEPLOYMENT
    destination.parent.mkdir(parents=True, exist_ok=True)
    torch.save(exported, destination)
    reloaded = torch.load(destination, map_location="cpu", weights_only=True)
    if reloaded["model"].keys() != saved["model"].keys() or any(
            not torch.equal(value, reloaded["model"][name]) for name, value in saved["model"].items()):
        raise ValueError("Inference export changed a model tensor")
    del saved

    model, _, meta = load(destination, ROOT / TOKENIZER_DIR / "tokenizer.model", "cpu")
    processor = spm.SentencePieceProcessor(model_file=str(ROOT / TOKENIZER_DIR / "tokenizer.model"))
    content = processor.encode(TEXT * 8, out_type=int)
    bos, pad = processor.bos_id(), processor.pad_id()
    if (bos, pad) != (2, 0) or len(content) < 127:
        raise ValueError("Unexpected tokenizer specials or fixture text too short")
    cases = {"bos": [bos], "short": [bos, *content[:15]], "long": [bos, *content[:127]]}
    cases["right_pad"] = cases["short"] + [pad] * (128 - len(cases["short"]))
    cases["changed_future"] = cases["short"] + content[:128 - len(cases["short"])]
    arrays = {}
    with torch.inference_mode():
        for name, ids in cases.items():
            inputs = torch.tensor([ids], dtype=torch.int32)
            arrays[f"{name}_input_ids"] = inputs.numpy()
            arrays[f"{name}_logits"] = model(inputs.long()).numpy()
        cache, steps = model.new_cache(), []
        long_ids = torch.tensor([cases["long"]], dtype=torch.int32)
        for position in range(long_ids.shape[1]):
            steps.append(model.forward_cached(long_ids[:, position:position + 1], cache))
        arrays["long_kv_step_logits"] = torch.cat(steps, dim=1).numpy()
    worst = float(np.abs(arrays["long_kv_step_logits"] - arrays["long_logits"]).max())
    if worst > 3e-4:
        raise ValueError(f"forward_cached differs from forward by {worst}")
    reference = payload / REFERENCE
    reference.mkdir(parents=True)
    np.savez_compressed(reference / "logits.npz", **arrays)
    write_json(reference / "reference.json", {
        "format": "vimeml_zh_v1_mac_preparation_reference_v1",
        "release_name": "vimeml-tiny-zh-v1", "checkpoint_step": meta["step"],
        "checkpoint_sha256": CHECKPOINT_SHA256, "deployment_checkpoint_sha256": file_sha(destination),
        "tokenizer_sha256": TOKENIZER_SHA256, "precision": "fp32", "device": "cpu", "torch_version": torch.__version__,
        "cases": {name: {"length": len(ids),
                         "valid_prefix": len(cases["short"]) if name in {"right_pad", "changed_future"} else len(ids)}
                  for name, ids in cases.items()},
        "kv_steps": {"case": "long", "array": "long_kv_step_logits", "max_abs_vs_stateless": worst,
                     "method": "TinyGPTZh13.forward_cached one token at a time from an empty cache"},
        "policy": "PyTorch preparation fixtures; not Core ML validation, production quality or device timing.",
    })
    return {"release_name": "vimeml-tiny-zh-v1", "arm": "12k-fresh-s1", "checkpoint_step": meta["step"],
            "source_checkpoint": CHECKPOINT, "source_checkpoint_sha256": CHECKPOINT_SHA256,
            "deployment_checkpoint": DEPLOYMENT, "deployment_checkpoint_sha256": file_sha(destination),
            "tokenizer": TOKENIZER_DIR, "tokenizer_sha256": TOKENIZER_SHA256,
            "model_config": exported["model_config"], "parameter_count": model.parameter_count(),
            "kv_cache_shape": [exported["model_config"]["n_layers"], 1, 6, 128, 64], "optimizer_included": False}


def prepare_benchmark(payload):
    items = [i for i in read_jsonl(ROOT / FROZEN / "items.jsonl") if i["split"] == "development"]
    ids = {i["id"] for i in items}
    if len(ids) != 1456:
        raise ValueError("Expected 1,456 development items")
    for name in FILTERED:
        rows = items if name == "items.jsonl" else [r for r in read_jsonl(ROOT / FROZEN / name) if r["id"] in ids]
        write_jsonl(payload / FROZEN / name, rows)
    source = json.loads((ROOT / FROZEN / "manifest.json").read_text(encoding="utf-8"))
    write_json(payload / FROZEN / "manifest.json", {
        "version": source["version"], "split_included": "development only", "labels": source["labels"],
        "reading_of_谁": source["reading_of_谁"], "items": len(items),
        "abbrev_items": sum(i["keys_abbrev"] is not None for i in items),
        "strata": {k: v for k, v in source["strata"].items() if "/development/" in k}})
    (payload / SCORES).mkdir(parents=True)
    for track in ("full", "abbrev"):
        if {r["id"] for r in read_jsonl(ROOT / SCORES / f"{track}.jsonl")} - ids:
            raise ValueError("Reference scores contain non-development items")
        for name in (f"{track}.jsonl", f"{track}.meta.json"):
            shutil.copy2(ROOT / SCORES / name, payload / SCORES / name)
    for directory in RESULTS:
        summary = json.loads((ROOT / directory / "summary.json").read_text(encoding="utf-8"))
        if summary["split"] != "development":
            raise ValueError(f"{directory} is not a development result")
        (payload / directory).mkdir(parents=True)
        for name in ("items.jsonl", "summary.json"):
            shutil.copy2(ROOT / directory / name, payload / directory / name)
    for track, expected in EXPECTED_TOP1.items():
        summary = json.loads((ROOT / RESULTS[0 if track == "full" else 1] / "summary.json").read_text(encoding="utf-8"))
        if summary["all"]["top1"] != expected:
            raise ValueError(f"v1 {track} development Top-1 is not {expected}")


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    output = args.output.resolve()
    if output.exists():
        parser.error("Output exists; choose a new handoff version.")
    if git("status", "--porcelain", "--untracked-files=no"):
        parser.error("Commit tracked source changes before creating the portable Git bundle.")
    if git("branch", "--show-current") != "main" or git("rev-parse", "HEAD") != git("rev-parse", "origin/main"):
        parser.error("Package the pushed main branch.")
    from vimeml.training.data import file_sha

    maintenance = ROOT / "outputs/maintenance"
    maintenance.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="mac-zh-v1-package-", dir=maintenance) as temporary:
        stage = Path(temporary)
        payload = stage / "payload"
        model = prepare_weights(payload)
        prepare_benchmark(payload)
        for relative in COPIED:
            target = payload / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(ROOT / relative, target)
        for path in payload.rglob("*"):
            if path.is_file() and (path.is_symlink() or path.name.startswith(".env") or path.suffix == ".pyc"):
                raise ValueError(f"Unexpected input: {path}")
        subprocess.run(["git", "bundle", "create", str(stage / "repository.bundle"), "main"], cwd=ROOT, check=True)
        subprocess.run(["git", "bundle", "verify", str(stage / "repository.bundle")], cwd=ROOT, check=True,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        manifest = {
            "format": "vimeml_windows_to_mac_zh_v1",
            "source_commit": git("rev-parse", "HEAD"),
            "source_branch": "main",
            "mac_branch": "main",
            "checkout_name": "VimeML-zh-v1-coreml",
            "handoff_record": "outputs/deployment/mac-zh-v1-handoff.json",
            "remote": git("remote", "get-url", "origin"),
            "model": model,
            "reference": REFERENCE,
            "benchmark": {"version": "ime-standard-zh-v1-r1", "split_included": "development",
                          "frozen": FROZEN, "reference_scores": SCORES,
                          "reference_top1": EXPECTED_TOP1, "rime_top1": {"full": 1001, "abbrev": 142}},
            "blind_included": False,
            "coreml_conversion_done": False,
            "payload_files": [
                {"path": p.relative_to(payload).as_posix(), "bytes": p.stat().st_size, "sha256": file_sha(p)}
                for p in sorted(payload.rglob("*")) if p.is_file()
            ],
            "verification": "Checkpoint and tokenizer SHA256; exact tensor equality; KV vs stateless logits; "
                            "development-only filter; reference Top-1; Git bundle verify; per-file SHA256; ZIP CRC.",
        }
        write_json(stage / "handoff-manifest.json", manifest)
        shutil.copy2(ROOT / "scripts/deployment/setup_mac.py", stage / "setup_mac.py")
        shutil.copy2(ROOT / "docs/reference/mac-zh-v1-coreml.md", stage / "README.md")
        output.parent.mkdir(parents=True, exist_ok=True)
        expected = {}
        with zipfile.ZipFile(output, "x", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
            for path in sorted(stage.rglob("*")):
                if path.is_file():
                    name = path.relative_to(stage).as_posix()
                    archive.write(path, name)
                    expected[name] = path.stat().st_size
        with zipfile.ZipFile(output) as archive:
            if {p.filename: p.file_size for p in archive.infolist()} != expected or archive.testzip() is not None:
                raise ValueError("ZIP directory or CRC differs from packaged files.")
        summary = {"zip": output.name, "bytes": output.stat().st_size, "sha256": file_sha(output),
                   "file_count": len(expected), "source_commit": manifest["source_commit"], "model": model}
        write_json(output.with_suffix(".json"), summary)
        print(json.dumps({k: summary[k] for k in ("zip", "bytes", "sha256", "file_count", "source_commit")},
                         ensure_ascii=False))


if __name__ == "__main__":
    main()