Working CoreML LUT4 input_embeds variant (86.9% on VITW)

inference_asr.py

@@ -0,0 +1,216 @@
+"""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))
+
+    # 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())