""" generate.py - Clean MLX inference for NVIDIA Canary-Qwen-2.5B Strategy: 1. Load base Qwen/Qwen3-1.7B weights (the actual NeMo pretrained_llm) 2. Merge NeMo LoRA deltas on top (q_proj, v_proj) 3. Load FastConformer encoder + audio projection 4. Run inference with the exact NeMo prompt format """ import sys import json import numpy as np import mlx.core as mx import mlx.nn as nn import librosa from mlx_lm.utils import load_model as mlx_load_model from mlx_lm.models.cache import make_prompt_cache from transformers import AutoTokenizer # ── paths ────────────────────────────────────────────────────────────────── CANARY_CKPT = "/Users/santosh/Desktop/ASR-Playground/canary_weights/model.safetensors" ENC_WEIGHTS = "/Users/santosh/Desktop/ASR-Playground/mlx_canary/mlx_canary_encoder.safetensors" QWEN_BASE_DIR = None # filled in at runtime SAMPLE_RATE = 16_000 N_MELS = 128 LORA_SCALE = 256 / 128 # alpha / r # ── audio feature extraction ──────────────────────────────────────────────── def extract_mel(audio_path: str) -> mx.array: y, _ = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True) # Pre-emphasis y = np.append(y[0], y[1:] - 0.97 * y[:-1]) # STFT matching NeMo config: n_fft=512, hop=160, win_length=400, center=True stft = librosa.stft(y, n_fft=512, hop_length=160, win_length=400, window="hann", center=True) power = np.abs(stft) ** 2 mel_fb = librosa.filters.mel(sr=SAMPLE_RATE, n_fft=512, n_mels=N_MELS, fmin=0.0, fmax=8000.0, norm="slaney", htk=False) with np.errstate(all="ignore"): mel = mel_fb @ power log_mel = np.log(np.maximum(mel, 1e-5)) # per-feature normalise (NeMo "per_feature") mean = log_mel.mean(axis=1, keepdims=True) std = log_mel.std(axis=1, keepdims=True) + 1e-5 log_mel = (log_mel - mean) / std # (features, time) -> (1, time, features) return mx.array(log_mel.T[None]) # ── LLM setup ─────────────────────────────────────────────────────────────── def load_llm(qwen_base_dir: str): """Load Qwen3-1.7B base via mlx_lm, then overlay NeMo LoRA-merged weights.""" from pathlib import Path # 1. Use mlx_lm's own load_model - handles qwen3, sanitize, etc. correctly model, _ = mlx_load_model(Path(qwen_base_dir)) # 2. Load NeMo checkpoint and overlay LoRA-merged q/v + fine-tuned layers nemo = mx.load(CANARY_CKPT) overlay: dict = {} for k, v in nemo.items(): if not k.startswith("llm.base_model.model.model."): continue new_k = k.replace("llm.base_model.model.model.", "model.") if ".base_layer.weight" in new_k: prefix = new_k.replace(".base_layer.weight", "") lora_a = nemo.get(k.replace(".base_layer.weight", ".lora_A.default.weight")) lora_b = nemo.get(k.replace(".base_layer.weight", ".lora_B.default.weight")) if lora_a is not None and lora_b is not None: overlay[f"{prefix}.weight"] = v + (lora_b @ lora_a) * LORA_SCALE else: overlay[f"{prefix}.weight"] = v continue if ".lora_A." in new_k or ".lora_B." in new_k: continue overlay[new_k] = v # Also load embed_tokens from the NeMo checkpoint (it stores the full updated embeddings) if "embed_tokens.weight" in nemo: overlay["model.embed_tokens.weight"] = nemo["embed_tokens.weight"] # Evaluate the merged LoRA weight graphs to materialize them as static tensors in GPU RAM mx.eval(list(overlay.values())) model.load_weights(list(overlay.items()), strict=False) # 3. Zero out any uninitialised QKV biases for layer in model.model.layers: for proj_name in ("q_proj", "k_proj", "v_proj"): proj = getattr(layer.self_attn, proj_name, None) if proj is not None and hasattr(proj, "bias"): proj.bias = mx.zeros_like(proj.bias) model.eval() return model def get_embed_tokens(model): """Safely get the embed_tokens layer regardless of qwen2 vs qwen3 model wrapper.""" if hasattr(model, 'model') and hasattr(model.model, 'embed_tokens'): return model.model.embed_tokens raise AttributeError("Cannot find embed_tokens on model") # ── encoder setup ──────────────────────────────────────────────────────────── def load_encoder(): import sys sys.path.insert(0, "/Users/santosh/Desktop/ASR-Playground/mlx_canary") from model import FastConformerEncoder, CanaryModel encoder = FastConformerEncoder( feat_in=128, n_layers=32, d_model=1024, ff_expansion_factor=4, n_heads=8, conv_kernel_size=9, dropout=0.0 ) canary = CanaryModel(encoder=encoder, llm_dim=2048) canary.load_weights(ENC_WEIGHTS, strict=False) canary.eval() return canary # ── main inference ─────────────────────────────────────────────────────────── def transcribe(audio_path: str, qwen_base_dir: str, use_ref_mel: bool = False) -> str: # Audio features if use_ref_mel: log_mel = np.load( "/Users/santosh/Desktop/ASR-Playground/mlx_canary/nemo_mel.npy" ) print(f"Mel shape: {log_mel.shape}") audio_features = mx.array(log_mel) else: audio_features = extract_mel(audio_path) print(f"Mel shape: {audio_features.shape}") # Encode audio canary = load_encoder() audio_embeds, _ = canary.encode_audio(audio_features) audio_std = mx.std(audio_embeds).item() print(f"Audio embeds: {audio_embeds.shape} mean={mx.mean(audio_embeds).item():.4f} std={audio_std:.4f}") # LLM llm = load_llm(qwen_base_dir) tokenizer = AutoTokenizer.from_pretrained(qwen_base_dir) embed_tokens = get_embed_tokens(llm) # Build prompt matching what NeMo's QwenPromptFormatter produces. # We use the exact tokens from the reference PyTorch SALM pipeline: # t1: "<|im_start|>user\nTranscribe the following: " # t2: "<|im_end|>\n<|im_start|>assistant\n" t1 = [151644, 872, 198, 3167, 3114, 279, 2701, 25, 220] t2 = [151645, 198, 151644, 77091, 198] print(f"Prompt t1 tokens ({len(t1)}): {t1}") print(f"Prompt t2 tokens ({len(t2)}): {t2}") emb1 = embed_tokens(mx.array([t1])) emb2 = embed_tokens(mx.array([t2])) text_std = mx.std(emb1).item() print(f"Text embeds (emb1): mean={mx.mean(emb1).item():.4f} std={text_std:.4f}") # No scaling is needed because the Conformer encoder is now corrected and matches the reference range # Cast audio embeddings to the LLM weight precision (typically bfloat16) to prevent MLX from # upcasting the entire model computation and KV Cache to float32, which disables GPU Tensor Core acceleration. scaled_audio_embeds = audio_embeds.astype(embed_tokens.weight.dtype) full_embeds = mx.concatenate([emb1, scaled_audio_embeds, emb2], axis=1) mx.eval(full_embeds) print(f"Full embed shape: {full_embeds.shape} dtype: {full_embeds.dtype}") cache = make_prompt_cache(llm) dummy = mx.zeros((1, full_embeds.shape[1]), dtype=mx.int32) logits = llm(dummy, cache=cache, input_embeddings=full_embeds) y = mx.argmax(logits[:, -1, :], axis=-1) mx.async_eval(y) def step(token): logits = llm(token, cache=cache) return mx.argmax(logits[:, -1, :], axis=-1) generated = [] im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>") # Evaluate first token mx.eval(y) token_id = y.item() print(f"First token: {tokenizer.decode([token_id])} ({token_id})") for _ in range(200): if token_id in (tokenizer.eos_token_id, im_end_id): break generated.append(token_id) # Enqueue the next token computation on GPU asynchronously next_y = step(mx.array([[token_id]])) mx.async_eval(next_y) # Get token value (blocks CPU while next_y is computed in parallel on GPU) token_id = next_y.item() return tokenizer.decode(generated) if __name__ == "__main__": import os, glob audio = sys.argv[1] if len(sys.argv) > 1 else "/Users/santosh/Desktop/ASR-Playground/test.wav" use_ref = "--ref" in sys.argv # Find the downloaded Qwen2-1.5B-Instruct model candidates = glob.glob(os.path.expanduser( "~/.cache/huggingface/hub/models--Qwen--Qwen3-1.7B/snapshots/*/")) if not candidates: print("ERROR: Qwen3-1.7B not found. Run:") print(" python -c \"from huggingface_hub import snapshot_download; snapshot_download('Qwen/Qwen3-1.7B')\"") sys.exit(1) qwen_dir = candidates[0].rstrip("/") print(f"Using Qwen base: {qwen_dir}") result = transcribe(audio, qwen_dir, use_ref_mel=use_ref) print(f"\n=== Transcription ===\n{result}\n=====================")