Upload generate.py with huggingface_hub

e3f9d15 verified 5 months ago

9.41 kB

"""
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=====================")

	"""
	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=====================")