| """End-to-end Mega-ASR pipeline on CoreML: |
| ONNX audio encoder + CoreML LUT-4 LLM (input_embeds variant) + bench. |
| |
| The CoreML LLM is single-token-step (ANE-friendly). For each token in the |
| prompt we feed (inputs_embeds[t], current_pos=t) to populate the KV cache; |
| then we greedy-decode by feeding one token at a time. |
| """ |
| from __future__ import annotations |
| import argparse |
| import json |
| import re |
| import sys |
| import warnings |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| warnings.filterwarnings("ignore") |
|
|
| REFERENCES = { |
| "noise": "I usually take the quieter road home because the main street gets crowded after work.", |
| "far_field": "Please remind me to print the forms before we leave for the appointment tomorrow.", |
| "obstructed": "I forgot my charger at home, so I need to find an outlet before the meeting starts.", |
| "distortion": "The new coffee machine is simple, but everyone keeps forgetting where the filters are stored.", |
| "recording": "Can you check whether the train still stops at the downtown station after eight tonight?", |
| "echo": "I need to return these shoes because the size feels fine standing up but terrible while walking.", |
| "dropout": "My aunt is learning video calls, and she gets excited whenever the picture actually works.", |
| "mixed": "My sister is bringing dinner over later, so we do not need to cook tonight.", |
| } |
| _NORM_RE = re.compile(r"[^a-z0-9\s]") |
|
|
|
|
| def normalize(t): |
| if "<asr_text>" in t: |
| t = t.split("<asr_text>", 1)[1] |
| return re.sub(r"\s+", " ", _NORM_RE.sub(" ", t.lower())).strip() |
|
|
|
|
| def wer(ref, hyp): |
| r = ref.split(); h = hyp.split() |
| if not r: return (1.0 if h else 0.0, len(h), 0) |
| d = np.zeros((len(r) + 1, len(h) + 1), dtype=np.int32) |
| for i in range(len(r) + 1): d[i, 0] = i |
| for j in range(len(h) + 1): d[0, j] = j |
| for i in range(1, len(r) + 1): |
| for j in range(1, len(h) + 1): |
| d[i, j] = min(d[i-1, j] + 1, d[i, j-1] + 1, |
| d[i-1, j-1] + (0 if r[i-1] == h[j-1] else 1)) |
| return d[len(r), len(h)] / max(len(r), 1), int(d[len(r), len(h)]), len(r) |
|
|
|
|
| def color(p, s): |
| if p >= 70: return f"\033[92m{s}\033[0m" |
| if p >= 50: return f"\033[93m{s}\033[0m" |
| if p >= 25: return f"\033[33m{s}\033[0m" |
| return f"\033[91m{s}\033[0m" |
|
|
|
|
| def build_prompt_ids(tok, audio_pad_count, audio_pad_id=151676): |
| prompt = ( |
| "<|im_start|>system\n<|im_end|>\n" |
| "<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|><|im_end|>\n" |
| "<|im_start|>assistant\n" |
| "language English<asr_text>" |
| ) |
| ids = tok.encode(prompt, add_special_tokens=False) |
| pos = ids.index(audio_pad_id) |
| return ids[:pos] + [audio_pad_id] * audio_pad_count + ids[pos + 1:] |
|
|
|
|
| def causal_mask_at(cur, ctx, neg_inf=-1e4): |
| """Build (1,1,1,ctx) mask: positions > cur get -inf, others 0.""" |
| m = np.zeros((1, 1, 1, ctx), dtype=np.float16) |
| if cur + 1 < ctx: |
| m[0, 0, 0, cur + 1:] = neg_inf |
| return m |
|
|
|
|
| def update_mask_at(cur, ctx): |
| """(1,1,ctx,1) — 1.0 at the current position, 0 elsewhere. Used for KV cache write.""" |
| m = np.zeros((1, 1, ctx, 1), dtype=np.float16) |
| m[0, 0, cur, 0] = 1.0 |
| return m |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--mlpackage", default="models/coreml/mega-asr-llm-embeds_lut4.mlpackage", type=Path) |
| ap.add_argument("--encoder-path", default="models/mega-asr-onnx-hf/onnx/audio_encoder_fp32.onnx", type=Path) |
| ap.add_argument("--examples-dir", default="models/mega-asr-onnx-hf/examples", type=Path) |
| ap.add_argument("--qwen-asr-dir", default="models/mega-asr-hf/Qwen3-ASR-1.7B", type=Path) |
| ap.add_argument("--max-new-tokens", type=int, default=80) |
| ap.add_argument("--context-length", type=int, default=512) |
| ap.add_argument("--compute-unit", default="CPU_AND_NE", choices=["CPU_ONLY", "CPU_AND_NE", "ALL"]) |
| args = ap.parse_args() |
|
|
| import soundfile as sf |
| import onnxruntime as ort |
| import coremltools as ct |
| from transformers import AutoFeatureExtractor, AutoTokenizer |
|
|
| print(f"Loading CoreML mlpackage ({args.compute_unit}) ...") |
| cu = getattr(ct.ComputeUnit, args.compute_unit) |
| mlm = ct.models.MLModel(str(args.mlpackage), compute_units=cu) |
|
|
| print(f"Loading ONNX encoder ...") |
| enc = ort.InferenceSession(str(args.encoder_path), providers=["CPUExecutionProvider"]) |
| feat = AutoFeatureExtractor.from_pretrained(str(args.qwen_asr_dir)) |
| tok = AutoTokenizer.from_pretrained(str(args.qwen_asr_dir)) |
|
|
| |
| import safetensors.torch as st |
| import torch |
| print("Loading embed_tokens (bf16 → fp32) ...") |
| |
| idx = json.load(open(args.qwen_asr_dir / "model.safetensors.index.json")) |
| embed_key = "thinker.model.embed_tokens.weight" |
| shard = idx["weight_map"][embed_key] |
| embed_w = st.load_file(str(args.qwen_asr_dir / shard))[embed_key] |
| embed_w = embed_w.to(torch.float32).numpy() |
| HIDDEN = embed_w.shape[1] |
| print(f" embed_w shape: {embed_w.shape}") |
|
|
| AUDIO_PAD = 151676 |
| EOS = 151645 |
| CTX = args.context_length |
|
|
| total_wer = 0.0; total_edits = 0; total_words = 0; n = 0 |
| for name in sorted(REFERENCES): |
| wav_path = args.examples_dir / f"{name}.wav" |
| if not wav_path.exists(): |
| print(f" skip {name} (missing)"); continue |
| audio, sr = sf.read(str(wav_path)) |
| if audio.ndim > 1: audio = audio.mean(axis=1) |
| if sr != 16000: |
| import librosa |
| audio = librosa.resample(audio.astype(np.float32), orig_sr=sr, target_sr=16000) |
| f = feat(audio, sampling_rate=16000, return_tensors="np", return_attention_mask=False) |
| mel = f["input_features"] |
| T_mel = mel.shape[-1] |
| if T_mel > 3000: mel = mel[..., :3000]; T_mel = 3000 |
| mel = np.pad(mel, ((0, 0), (0, 0), (0, 3000 - T_mel)), constant_values=0).astype(np.float32) |
| audio_embeds = enc.run(["audio_embeds"], {"mel": mel})[0] |
| real_chunks = (T_mel + 99) // 100 |
| last_chunk = T_mel - (real_chunks - 1) * 100 |
| real_frames = (real_chunks - 1) * 13 + (last_chunk + 7) // 8 |
| audio_embeds = audio_embeds[0, :real_frames] |
|
|
| |
| prompt_ids = build_prompt_ids(tok, real_frames) |
| L = len(prompt_ids) |
| if L > CTX - args.max_new_tokens: |
| print(f" skip {name} (L={L} too long for ctx={CTX})"); continue |
|
|
| |
| token_embeds = np.zeros((L, HIDDEN), dtype=np.float32) |
| ai = 0 |
| for i, t in enumerate(prompt_ids): |
| if t == AUDIO_PAD: |
| token_embeds[i] = audio_embeds[ai]; ai += 1 |
| else: |
| token_embeds[i] = embed_w[t] |
| token_embeds = token_embeds.astype(np.float16) |
|
|
| |
| state = mlm.make_state() |
| for i in range(L): |
| feeds = { |
| "inputs_embeds": token_embeds[i:i+1].reshape(1, 1, HIDDEN), |
| "position_ids": np.array([i], dtype=np.int32), |
| "causal_mask": causal_mask_at(i, CTX), |
| "current_pos": np.array([i], dtype=np.int32), |
| "update_mask": update_mask_at(i, CTX), |
| } |
| out = mlm.predict(feeds, state=state) |
| |
| logits = np.concatenate([out[f"logits{k}"][0, 0] for k in range(1, 17)]) |
| nid = int(np.argmax(logits)) |
| gen = [nid] |
|
|
| |
| for step in range(args.max_new_tokens - 1): |
| if nid == EOS: break |
| cur = L + step |
| if cur >= CTX: break |
| emb = embed_w[nid].astype(np.float16) |
| feeds = { |
| "inputs_embeds": emb.reshape(1, 1, HIDDEN), |
| "position_ids": np.array([cur], dtype=np.int32), |
| "causal_mask": causal_mask_at(cur, CTX), |
| "current_pos": np.array([cur], dtype=np.int32), |
| "update_mask": update_mask_at(cur, CTX), |
| } |
| out = mlm.predict(feeds, state=state) |
| logits = np.concatenate([out[f"logits{k}"][0, 0] for k in range(1, 17)]) |
| nid = int(np.argmax(logits)) |
| gen.append(nid) |
| |
| if gen and gen[-1] == EOS: gen = gen[:-1] |
| hyp_text = tok.decode(gen, skip_special_tokens=True) |
|
|
| ref = normalize(REFERENCES[name]) |
| hyp = normalize(hyp_text) |
| w, ed, words = wer(ref, hyp) |
| agree = max(0.0, 1.0 - w) * 100 |
| total_wer += w; total_edits += ed; total_words += words; n += 1 |
| print(f"\n[{color(agree, name.ljust(10))}] WER={w*100:5.1f}% agree={color(agree, f'{agree:5.1f}%')}") |
| print(f" REF: {ref}") |
| print(f" HYP: {hyp}") |
| avg = (1 - total_wer / n) * 100 if n else 0 |
| print(f"\n{color(avg, f'=== AVERAGE: agreement {avg:.1f}% WER {total_edits/total_words*100:.1f}% ({total_edits}/{total_words}) ===')}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
