"""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 "" in t: t = t.split("", 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" ) 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)) # Embed table from the HF model (for text tokens; audio_pad slots use audio_embeds) import safetensors.torch as st import torch print("Loading embed_tokens (bf16 → fp32) ...") # Find embed_tokens.weight from the original Qwen3-ASR safetensors 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() # (151936, 2048) 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] # (F, 2048) # Build prompt tokens + per-position embeddings 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 # Per-token embeddings: lookup for text, scatter audio_embeds at audio_pad slots 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) # Run prefill: feed each prompt token one at a time 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) # Argmax of last logit logits = np.concatenate([out[f"logits{k}"][0, 0] for k in range(1, 17)]) nid = int(np.argmax(logits)) gen = [nid] # Decode step-by-step 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) # Strip trailing eos 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())