| """Prepare vimeml-tiny-zh-v1 INT8 client resources and native fixtures for the Vime test host. |
| |
| Resources are named so that Japanese resources are unaffected: |
| TinyChineseV1INT8.mlmodelc, TinyChineseV1KVChunkINT8.mlmodelc, |
| VimeChineseTokenizerV1.model, VimeLMManifestZhV1.json (+ VimeLMFixturesZhV1.json beside Resources/). |
| |
| Fixture texts are written here; none comes from the benchmark. Device runs belong to the Vime repository. |
| """ |
| import argparse |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "src")) |
|
|
| import numpy as np |
| import torch |
|
|
| from vimeml.benchmarks.zh13_lm_scores import score_item |
| from vimeml.deployment import v2, zh_v1 as zh |
| from vimeml.deployment.bundle import fresh_directory, read_json, write_json |
| from vimeml.training.data import file_sha |
|
|
| TAG = "ZhV1" |
| TOKENIZATION = ["你好，今天天气不错。", "我们明天下午三点在地铁站见面吧。", "这家餐厅的火锅很好吃，推荐给你。", |
| "Vime 输入法支持拼音。", "一二三四五六七八九十"] |
| SCORING = [("我们晚上去吃", ["火锅", "活过", "伙夫", "或"]), |
| ("明天上午我要去", ["开会", "看会", "开灰", "凯"]), |
| ("这本书写得", ["很好", "恨好", "很号"]), |
| ("", ["你好", "拟好", "泥嚎"])] |
|
|
|
|
| def compile_package(package, destination): |
| subprocess.run(["xcrun", "coremlcompiler", "compile", str(package / "model.mlpackage"), str(destination)], |
| check=True, capture_output=True) |
| compiled = destination / "model.mlmodelc" |
| if not (compiled / "coremldata.bin").is_file(): |
| raise ValueError("Missing compiled model.") |
| return compiled |
|
|
|
|
| @torch.inference_mode() |
| def package(bundle, stateless, chunk, alignment, output): |
| if output.exists(): |
| raise ValueError("Resource output exists.") |
| stateless_manifest = zh.verify_package(stateless) |
| chunk_manifest = zh.verify_package(chunk, zh.CHUNK_FORMAT) |
| for manifest in (stateless_manifest, chunk_manifest): |
| if manifest.get("recipe") != zh.DEPLOYED_RECIPE or manifest["minimum_ios"] != 18: |
| raise ValueError("Expected the deployed INT8 block32 FP32-compute iOS18 recipe.") |
| gate = read_json(alignment) |
| if (gate.get("format") != zh.CHUNK_ALIGNMENT_FORMAT or not gate.get("passed") |
| or gate.get("comparison") != "same-weight stateless Core ML" |
| or gate["coreml_manifest_sha256"] != file_sha(chunk / "manifest.json") |
| or gate["baseline_manifest_sha256"] != file_sha(stateless / "manifest.json")): |
| raise ValueError("Passed same-weight chunk-KV/stateless INT8 comparison required.") |
| lm = zh.CoreMLLMZh(bundle, stateless) |
| runtime = zh.ChunkRuntimeZh(bundle, chunk) |
| bos = lm.special["bos"] |
| fixtures = {"format": "vime_client_fixtures_zh_v1", "model": lm.metadata, "tokenization": [], "scoring": [], |
| "kv_chunk_steps": []} |
| for text in TOKENIZATION: |
| ids = lm.processor.encode(text) |
| fixtures["tokenization"].append({"text": text, "ids": ids, "decoded": lm.processor.decode(ids)}) |
| deltas = [] |
| for context, candidates in SCORING: |
| scored = score_item(lm.model, lm.processor, context, candidates) |
| common = scored["common_prefix_tokens_including_bos"] |
| fixtures["scoring"].append({"context": context, "candidates": candidates, |
| "common_prefix_tokens_including_bos": common, "sums": scored["scores"], |
| "scored_tokens": scored["scored_tokens"]}) |
| sequences = [[bos, *lm.processor.encode(context + c)] for c in candidates] |
| prefix_ids = sequences[0][:common] |
| |
| prefix_logits, snapshot = runtime.run(prefix_ids) |
| for text, sequence, expected in zip(candidates, sequences, scored["scores"], strict=True): |
| if sequence[:common] != prefix_ids: |
| raise ValueError("Common prefix differs from the scored boundary.") |
| rows, state = [prefix_logits[0, -1]], snapshot |
| for position in range(common, len(sequence) - 1): |
| logits, state = runtime.run([sequence[position]], state, position) |
| rows.append(logits[0, -1]) |
| logp = torch.from_numpy(np.stack(rows)).log_softmax(-1) |
| targets = sequence[common:] |
| values = [float(logp[i, t]) for i, t in enumerate(targets)] |
| deltas.append(abs(sum(values) - expected)) |
| fixtures["kv_chunk_steps"].append({"context": context, "text": text, "prefix_ids": prefix_ids, |
| "prefix_chunk_calls": 1, "decode_ids": sequence[common:-1], "targets": targets, |
| "target_log_probabilities": values, "suffix_sum": sum(values), "stateless_sum": expected}) |
| fresh_directory(output) |
| resources = output / "Resources" |
| resources.mkdir() |
| staging = output / "compile" |
| compiled = {"stateless": compile_package(stateless, staging / "stateless"), |
| "kv": compile_package(chunk, staging / "kv")} |
| names = {"stateless": "TinyChineseV1INT8.mlmodelc", "kv": "TinyChineseV1KVChunkINT8.mlmodelc"} |
| for key, name in names.items(): |
| shutil.copytree(compiled[key], resources / name) |
| shutil.rmtree(staging) |
| shutil.copyfile(Path(bundle) / "tokenizer.model", resources / "VimeChineseTokenizerV1.model") |
| bundle_manifest = read_json(Path(bundle) / "manifest.json") |
| manifest = {"format": "vime_ios_lm_zh_v1", "architecture": zh.ARCHITECTURE, "release_name": zh.RELEASE, |
| "model_version": f"zh-v1-step{lm.metadata['checkpoint_step']}-int8-b32-v1", |
| "minimum_ios": 18, "vocab_size": zh.VOCAB, "n_layers": zh.LAYERS, "context_length": zh.CONTEXT, |
| "compute_units": "CPU_ONLY", "special_ids": lm.special, |
| "reserved_ids": bundle_manifest.get("reserved_ids", {}), |
| "tokenizer_sha256": lm.metadata["tokenizer_sha256"], "checkpoint_sha256": lm.metadata["checkpoint_sha256"], |
| "bundle_manifest_sha256": lm.metadata["bundle_manifest_sha256"], |
| "models": {key: {"resource": names[key], |
| "source_manifest_sha256": file_sha((stateless if key == "stateless" else chunk) / "manifest.json"), |
| "compiled_files_sha256": {n: e["sha256"] for n, e in |
| v2.inventory(resources / names[key], hashes=True).items()}} |
| for key in names}, |
| "stateless_interface": stateless_manifest["interface"], "kv_interface": chunk_manifest["interface"], |
| "cache_shape": chunk_manifest["cache_shape"], "cache_dtype": "float32", "kv_graph": "chunk", |
| "kind": "linear8_fp32_compute", "quantization": "weight-only symmetric INT8 block32; FP32 computation", |
| "strict_fp32_logits_alignment_passed": False, |
| "kv_vs_stateless_alignment_sha256": file_sha(alignment), |
| "kv_chunk_fixture_max_suffix_sum_delta": max(deltas), |
| "kv_chunk_fixture_within_1e-4": max(deltas) <= 1e-4, |
| "scoring": "BOS + joint context/candidate tokenization; logP sum from the first token differing from the context alone; no EOS", |
| "quality_review": "Test-host measurement candidate only; not shipped in Vime. See VimeML docs/reports/mac-*/zh-v1-coreml.md."} |
| write_json(resources / f"VimeLMManifest{TAG}.json", manifest) |
| write_json(output / f"VimeLMFixtures{TAG}.json", fixtures) |
| write_json(output / "manifest.json", {"format": "vimeml_zh_v1_ios_resources_v1", "model": lm.metadata, |
| "files": v2.inventory(output), |
| "note": "Resources/ is for the Vime test host. No client source snapshot or device run included."}) |
| print(f"KV chunk fixture max suffix-sum delta: {max(deltas):.3g}", flush=True) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| for name in ("bundle", "stateless", "chunk", "chunk-alignment", "output"): |
| parser.add_argument("--" + name, type=Path, required=True) |
| args = parser.parse_args() |
| torch.set_num_threads(2) |
| package(args.bundle, args.stateless, args.chunk, args.chunk_alignment, args.output) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
