onnx-community

/

Voxtral-Mini-3B-2507-ONNX

like 29

Follow

ONNX Community 2.09k

import numpy as np
import onnxruntime as ort
import soundfile as sf
import librosa
from tokenizers import Tokenizer

# --- 1. Audio preprocessing ---
audio, sr = sf.read("audio.wav")
if len(audio.shape) > 1:
    audio = np.mean(audio, axis=1)
if sr != 16000:
    audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
n_mels = 128
n_frames = 3000
mel = librosa.feature.melspectrogram(
    y=audio, sr=sr, n_fft=400, hop_length=160, n_mels=n_mels, fmin=0, fmax=8000
)
log_mel = np.log10(np.maximum(mel, 1e-10))
if log_mel.shape[1] < n_frames:
    pad_width = n_frames - log_mel.shape[1]
    log_mel = np.pad(log_mel, ((0, 0), (0, pad_width)), mode='constant', constant_values=0)
else:
    log_mel = log_mel[:, :n_frames]
log_mel = log_mel[np.newaxis, ...].astype(np.float32)  # (1, 128, 3000)

# --- 2. Load ONNX models and tokenizer ---
onnx_dir = "../voxtral_onnx_models/onnx"
tokenizer_path = "../voxtral_onnx_models/tokenizer.json"

encoder_sess = ort.InferenceSession(f"{onnx_dir}/audio_encoder.onnx")
decoder_sess = ort.InferenceSession(f"{onnx_dir}/decoder_model_merged.onnx", providers=["CPUExecutionProvider"])
embed_sess = ort.InferenceSession(f"{onnx_dir}/embed_tokens.onnx")
embed_input_name = embed_sess.get_inputs()[0].name

text_tokenizer = Tokenizer.from_file(tokenizer_path)

# --- 3. Encode audio to get audio embedding sequence (already projected) ---
encoder_input_name = encoder_sess.get_inputs()[0].name
audio_embeds = encoder_sess.run(None, {encoder_input_name: log_mel})[0]  # (1, N, 3072)
if audio_embeds.ndim == 2:
    audio_embeds = audio_embeds[np.newaxis, :, :]
num_audio_frames = audio_embeds.shape[1]

# --- 4. Build the exact prompt as discovered ---
bos_token_id = text_tokenizer.token_to_id("<s>")
eos_token_id = text_tokenizer.token_to_id("</s>")
inst_start_token_id = 3
begin_audio_token_id = 25
audio_token_placeholder_id = 24
inst_end_token_id = 4
lang_colon_token_id = 9909
en_token_id = 1058
sq_bracket_open_token_id = 1262
transcribe_token_id = 34

segment1_ids = np.array([bos_token_id, inst_start_token_id, begin_audio_token_id], dtype=np.int64)
audio_placeholder_ids_sequence = np.full(num_audio_frames, audio_token_placeholder_id, dtype=np.int64)
segment3_ids = np.array([
    inst_end_token_id,
    lang_colon_token_id,
    en_token_id,
    sq_bracket_open_token_id,
    transcribe_token_id
], dtype=np.int64)

prompt_ids = np.concatenate([segment1_ids, audio_placeholder_ids_sequence, segment3_ids])
prompt_ids_batch = prompt_ids[np.newaxis, :]

# --- 5. Get embeddings for the prompt (including [AUDIO] tokens) ---
all_prompt_embeds = embed_sess.run(None, {embed_input_name: prompt_ids_batch})[0]  # (1, total_prompt_len, 3072)

# --- 6. Replace [AUDIO] token embeddings with audio embeddings ---
audio_token_positions = np.where(prompt_ids == audio_token_placeholder_id)[0]
assert len(audio_token_positions) == num_audio_frames, "Mismatch between [AUDIO] tokens and audio frames!"
for i, pos in enumerate(audio_token_positions):
    all_prompt_embeds[0, pos, :] = audio_embeds[0, i, :]

inputs_embeds = all_prompt_embeds  # (1, total_prompt_len, 3072)

# --- 7. Prepare for generation ---
batch_size = 1
num_layers = 30
num_heads = 8
head_dim = 128
embed_dim = 3072

seq_len = inputs_embeds.shape[1]
attention_mask = np.ones((batch_size, seq_len), dtype=np.int64)
position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :]

past_shape = (batch_size, num_heads, 0, head_dim)
past_key_values = {}
for layer in range(num_layers):
    past_key_values[f"past_key_values.{layer}.key"] = np.empty(past_shape, dtype=np.float32)
    past_key_values[f"past_key_values.{layer}.value"] = np.empty(past_shape, dtype=np.float32)

# --- 8. Generation loop ---
max_new_tokens = 5
generated_ids = []
for step in range(max_new_tokens):
    decoder_inputs = {
        "inputs_embeds": inputs_embeds[:, -1:, :],
        "attention_mask": attention_mask,
        "position_ids": position_ids[:, -1:],
        **past_key_values
    }
    decoder_outputs = decoder_sess.run(None, decoder_inputs)
    logits = decoder_outputs[0]
    next_token_id = int(np.argmax(logits[0, -1]))
    if next_token_id == eos_token_id and step == 0:
        print("Model immediately generated EOS. Check prompt carefully.")
        break
    generated_ids.append(next_token_id)
    if next_token_id == eos_token_id:
        break
    present_key_values = {}
    offset = 1
    for layer in range(num_layers):
        present_key = decoder_outputs[offset + 2 * layer]
        present_value = decoder_outputs[offset + 2 * layer + 1]
        present_key_values[f"past_key_values.{layer}.key"] = present_key
        present_key_values[f"past_key_values.{layer}.value"] = present_value
    past_key_values = present_key_values
    next_token_ids_batch = np.array([[next_token_id]], dtype=np.int64)
    next_token_embed = embed_sess.run(None, {embed_input_name: next_token_ids_batch})[0]
    inputs_embeds = np.concatenate([inputs_embeds, next_token_embed], axis=1)
    attention_mask = np.concatenate([attention_mask, np.ones((batch_size, 1), dtype=np.int64)], axis=1)
    position_ids = np.concatenate([position_ids, np.array([[attention_mask.shape[1] - 1]], dtype=np.int64)], axis=1)

# --- 9. Decode tokens to text ---
print("\nGenerated transcription:")
print("=" * 80)
if not generated_ids or generated_ids[0] == eos_token_id:
    print("<NO TEXT GENERATED BEFORE EOS>")
else:
    if generated_ids and generated_ids[-1] == eos_token_id:
        transcribed_text = text_tokenizer.decode(generated_ids[:-1], skip_special_tokens=False)
    else:
        transcribed_text = text_tokenizer.decode(generated_ids, skip_special_tokens=False)
    print(transcribed_text)
print("=" * 80)

import os
import numpy as np
import onnxruntime as ort
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer
import soundfile as sf
import librosa

# --- Configuration ---
repo_id = "onnx-community/Voxtral-Mini-3B-2507-ONNX"
audio_file_path = "audioCoco.mp3"  # Path to your local audio file
max_generation_tokens = 200
eos_token_id = 2  # </s> token ID

# --- 1. Download models and tokenizer files ---
print(f"Downloading model files from {repo_id}...")
local_dir = snapshot_download(
    repo_id=repo_id,
    repo_type="model",
    allow_patterns=["onnx/*", "tokenizer.json", "config.json"]
)
onnx_dir = os.path.join(local_dir, "onnx")
print(f"Files downloaded to: {local_dir}")

# --- 2. Load the tokenizer ---
print("\nLoading tokenizer...")
tok = Tokenizer.from_file(os.path.join(local_dir, "tokenizer.json"))
bos_id, inst_id, baud_id, aud_id, einst_id = 1, 3, 25, 24, 4
print("Tokenizer loaded.")

# --- 3. Initialize ONNX Runtime sessions ---
print("\nInitializing ONNX Runtime sessions...")
ae_path = os.path.join(onnx_dir, "audio_encoder.onnx")
embed_path = os.path.join(onnx_dir, "embed_tokens.onnx")
dec_path = os.path.join(onnx_dir, "decoder_model_merged.onnx")
session_providers = ["CPUExecutionProvider"]
ae_sess = ort.InferenceSession(ae_path, providers=session_providers)
embed_sess = ort.InferenceSession(embed_path, providers=session_providers)
dec_sess = ort.InferenceSession(dec_path, providers=session_providers)
print("Sessions initialized.")

num_decoder_layers = sum(1 for i in dec_sess.get_inputs() if i.name.endswith(".key"))
print(f"Detected {num_decoder_layers} decoder layers.")

# --- 4. Audio Pre-processing Functions ---
def extract_mel_features_for_chunk(audio_chunk, sampling_rate=16000, n_fft=400, hop_length=160, n_mels=128, target_length=3000):
    """
    Processes a single audio chunk into a normalized log-mel spectrogram.
    This function expects the chunk to be <= 30 seconds.
    """
    # Pad or truncate the chunk to exactly 30 seconds
    target_samples = sampling_rate * 30
    audio_chunk = librosa.util.fix_length(audio_chunk, size=target_samples)

    mel_spec = librosa.feature.melspectrogram(
        y=audio_chunk, sr=sampling_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels
    )

    log_spec = np.log10(np.maximum(mel_spec, 1e-10))
    log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
    log_spec = (log_spec + 4.0) / 4.0

    if log_spec.shape[1] > target_length:
        log_spec = log_spec[:, :target_length]

    return log_spec.astype(np.float32)

def process_long_audio(audio_path, session, sampling_rate=16000):
    """
    Loads a long audio file, processes it in 30-second chunks,
    and concatenates the resulting audio embeddings.
    """
    print(f"Loading long audio from {audio_path}...")
    y, sr = sf.read(audio_path)
    if y.ndim > 1: y = y.mean(axis=1)
    if sr != sampling_rate:
        y = librosa.resample(y, orig_sr=sr, target_sr=sampling_rate)

    chunk_duration = 30
    chunk_samples = chunk_duration * sampling_rate
    num_chunks = int(np.ceil(len(y) / chunk_samples))
    all_audio_embeds = []

    print(f"Audio duration: {len(y)/sampling_rate:.2f}s. Processing in {num_chunks} chunk(s)...")
    for i in range(num_chunks):
        start_sample = i * chunk_samples
        end_sample = start_sample + chunk_samples
        chunk = y[start_sample:end_sample]

        # Get mel features for the current chunk (this function handles padding for the last chunk)
        mel_features = extract_mel_features_for_chunk(chunk, sampling_rate)
        audio_values = mel_features[None, :] # Add batch dim

        # Run the audio encoder for this chunk
        chunk_embeds_raw = session.run(None, {session.get_inputs()[0].name: audio_values})[0]
        all_audio_embeds.append(chunk_embeds_raw)
        print(f"  Processed chunk {i+1}/{num_chunks}")

    # Concatenate embeddings from all chunks
    concatenated_embeds = np.concatenate(all_audio_embeds, axis=0)
    return concatenated_embeds

# --- 5. Process Audio and Get Embeddings ---
print("\nProcessing audio...")
try:
    if not os.path.exists(audio_file_path):
        raise FileNotFoundError(f"Audio file '{audio_file_path}' not found.")
    
    audio_embeds_raw = process_long_audio(audio_file_path, ae_sess)
    batch_size = 1
    audio_output_frames = audio_embeds_raw.shape[0] // batch_size
    audio_embeds = audio_embeds_raw.reshape(batch_size, audio_output_frames, -1)
    print(f"\nSuccessfully processed long audio. Final embedding shape: {audio_embeds.shape}")

except Exception as e:
    print(f"An error occurred during audio processing: {e}. The script cannot continue.")
    exit()

# --- 6. Build the initial prompt ---
text_instruction_ids = tok.encode("Transcribe the audio.", add_special_tokens=False).ids
prompt_tokens = [bos_id, inst_id, baud_id] + [aud_id] * audio_output_frames + text_instruction_ids + [einst_id]
current_input_ids = np.array([prompt_tokens], dtype=np.int64)
initial_sequence_length = current_input_ids.shape[1]
print(f"\nInitial prompt sequence length: {initial_sequence_length}")

# --- 7. Run embed_tokens.onnx and splice audio embeddings ---
print("\nRunning embed_tokens.onnx and splicing...")
embed_inputs = {embed_sess.get_inputs()[0].name: current_input_ids}
inputs_embeds = embed_sess.run(None, embed_inputs)[0]
audio_splice_start_idx = 3
inputs_embeds[0, audio_splice_start_idx : audio_splice_start_idx + audio_output_frames, :] = audio_embeds[0]
print("Splicing complete.")

# --- 8. Generation Loop ---
print("\nStarting generation loop...")
generated_ids = []
past_key_values_cache = None
current_past_sequence_length = 0

for i in range(max_generation_tokens):
    dec_inputs = {}
    current_step_sequence_length = 1 if i > 0 else initial_sequence_length
    current_total_sequence_length = current_past_sequence_length + current_step_sequence_length

    if i == 0:  # Prefill
        dec_inputs["inputs_embeds"] = inputs_embeds
        dec_inputs["attention_mask"] = np.ones((batch_size, initial_sequence_length), dtype=np.int64)
        dec_inputs["position_ids"] = np.arange(initial_sequence_length, dtype=np.int64)[None, :]
        for l in range(num_decoder_layers):
            dec_inputs[f"past_key_values.{l}.key"] = np.zeros((batch_size, 8, 0, 128), dtype=np.float32)
            dec_inputs[f"past_key_values.{l}.value"] = np.zeros((batch_size, 8, 0, 128), dtype=np.float32)
    else:  # Decode
        next_token_input_id = np.array([[generated_ids[-1]]], dtype=np.int64)
        single_token_embeds = embed_sess.run(None, {"input_ids": next_token_input_id})[0]
        dec_inputs["inputs_embeds"] = single_token_embeds
        dec_inputs["attention_mask"] = np.ones((batch_size, current_total_sequence_length), dtype=np.int64)
        dec_inputs["position_ids"] = np.array([[current_past_sequence_length]], dtype=np.int64)
        for l in range(num_decoder_layers):
            dec_inputs[f"past_key_values.{l}.key"] = past_key_values_cache[l * 2]
            dec_inputs[f"past_key_values.{l}.value"] = past_key_values_cache[l * 2 + 1]

    outputs = dec_sess.run(None, dec_inputs)
    logits, past_key_values_cache = outputs[0], outputs[1:]
    next_token_id = int(np.argmax(logits[0, -1, :]))
    generated_ids.append(next_token_id)
    
    decoded_token_text = tok.decode([next_token_id], skip_special_tokens=True)
    print(f"Step {i+1}: ID={next_token_id:<5}, Decoded='{decoded_token_text}'", end="\r")

    current_past_sequence_length = current_total_sequence_length
    if next_token_id == eos_token_id:
        print("\n\nEnd-of-sequence token detected (</s>). Stopping generation.")
        break

print("\n\n--- Generation Complete ---")
final_generated_text = tok.decode(generated_ids, skip_special_tokens=True)
print(f"\nFinal generated text: \n{final_generated_text}")