"""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()