File size: 3,911 Bytes
64f7370 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python3
"""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] # iphone-asr/
OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else WORKSPACE / "dist/ASRModels.bundle"
BUNDLE_VERSION = "1.0.0"
SOURCES = {
# logical name -> (source path, dest file name)
"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":
# precompile so devices never pay the mlpackage->mlmodelc conversion
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})")
|