| |
| """Builds ASRModels.bundle for SpeechKit distribution. |
| |
| Collects models + assets from the iphone-asr workspace, precompiles |
| mlpackage -> mlmodelc (xcrun coremlcompiler), computes SHA-256 for every |
| asset, and writes a schema-v2 manifest.json. |
| |
| Usage: python3 make_asset_bundle.py [output_dir] |
| """ |
| import hashlib |
| import json |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| WORKSPACE = Path(__file__).resolve().parents[2] |
| OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else WORKSPACE / "dist/ASRModels.bundle" |
| BUNDLE_VERSION = "1.0.0" |
|
|
| SOURCES = { |
| |
| "tower": (WORKSPACE / "artifacts/audio_tower_multifunction_w8.mlpackage", "tower.mlpackage"), |
| "lm_prefill": (WORKSPACE / "artifacts/lm_shared/lm_cache_prefill_int4.onnx", "lm_prefill.onnx"), |
| "lm_decode": (WORKSPACE / "artifacts/lm_shared/lm_cache_decode_int4.onnx", "lm_decode.onnx"), |
| "lm_shared_data": (WORKSPACE / "artifacts/lm_shared/lm_shared_int4.data", "lm_shared_int4.data"), |
| "mask_gen": (WORKSPACE / "artifacts/mask_gen.onnx", "mask_gen.onnx"), |
| "asr_manifest": (WORKSPACE / "ios_assets/manifest.json", "asr_manifest.json"), |
| "hann_window": (WORKSPACE / "ios_assets/hann_window.bin", "hann_window.bin"), |
| "mel_filters": (WORKSPACE / "ios_assets/mel_filters.bin", "mel_filters.bin"), |
| "projector_norm_w": (WORKSPACE / "ios_assets/projector_norm_w.bin", "projector_norm_w.bin"), |
| "projector_norm_b": (WORKSPACE / "ios_assets/projector_norm_b.bin", "projector_norm_b.bin"), |
| "projector_lin_w": (WORKSPACE / "ios_assets/projector_lin_w.bin", "projector_lin_w.bin"), |
| "projector_lin_b": (WORKSPACE / "ios_assets/projector_lin_b.bin", "projector_lin_b.bin"), |
| "token_embedding": (WORKSPACE / "ios_assets/token_embedding_fp16.bin", "token_embedding_fp16.bin"), |
| "vocab_bytes": (WORKSPACE / "ios_assets/vocab_bytes.bin", "vocab_bytes.bin"), |
| "vocab_offsets": (WORKSPACE / "ios_assets/vocab_offsets.bin", "vocab_offsets.bin"), |
| } |
|
|
|
|
| def sha256_path(p: Path) -> str: |
| h = hashlib.sha256() |
| if p.is_dir(): |
| for f in sorted(p.rglob("*")): |
| if f.is_file(): |
| h.update(f.relative_to(p).as_posix().encode()) |
| h.update(f.read_bytes()) |
| else: |
| with p.open("rb") as fh: |
| for chunk in iter(lambda: fh.read(1 << 23), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def dir_size(p: Path) -> int: |
| if p.is_file(): |
| return p.stat().st_size |
| return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) |
|
|
|
|
| if OUT.exists(): |
| shutil.rmtree(OUT) |
| OUT.mkdir(parents=True) |
|
|
| assets = {} |
| for logical, (src, dest_name) in SOURCES.items(): |
| assert src.exists(), f"missing source: {src}" |
| dest = OUT / dest_name |
| if src.suffix == ".mlpackage": |
| |
| dest = OUT / (Path(dest_name).stem + ".mlmodelc") |
| subprocess.run(["xcrun", "coremlcompiler", "compile", str(src), str(OUT)], |
| check=True, capture_output=True) |
| compiled = OUT / (src.stem + ".mlmodelc") |
| if compiled != dest: |
| compiled.rename(dest) |
| elif src.is_dir(): |
| shutil.copytree(src, dest) |
| else: |
| shutil.copy2(src, dest) |
| assets[logical] = { |
| "file": dest.name, |
| "sha256": sha256_path(dest), |
| "sizeBytes": dir_size(dest), |
| } |
| print(f" {logical:18s} -> {dest.name} ({dir_size(dest)/1e6:.1f}MB)") |
|
|
| manifest = { |
| "schemaVersion": 2, |
| "bundleVersion": BUNDLE_VERSION, |
| "assets": assets, |
| "extra": {"model": "minillm-asr-zh", "builtBy": "make_asset_bundle.py"}, |
| } |
| (OUT / "manifest.json").write_text(json.dumps(manifest, indent=2)) |
| total = sum(a["sizeBytes"] for a in assets.values()) |
| print(f"\nbundle: {OUT} ({total/1e6:.0f}MB, schema v2, version {BUNDLE_VERSION})") |
|
|
