[model] support voxtral (speech-to-text) (#21635)
Co-authored-by: mengxiancheng03 <mengxiancheng03@kuaishou.com>
This commit is contained in:
Xiancheng Meng
co-authored by
mengxiancheng03
mengxiancheng03
parent
904bb476d8
commit
71544f0341
@@ -1333,6 +1333,7 @@ multimodal_model_archs = [
|
||||
"InternS1ForConditionalGeneration",
|
||||
"InternS1ProForConditionalGeneration",
|
||||
"Phi4MMForCausalLM",
|
||||
"VoxtralForConditionalGeneration",
|
||||
"WhisperForConditionalGeneration",
|
||||
"Step3VLForConditionalGeneration",
|
||||
"POINTSV15ChatModel",
|
||||
|
||||
@@ -0,0 +1,444 @@
|
||||
# Adapted from:
|
||||
# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/voxtral.py
|
||||
# https://huggingface.co/mistralai/Voxtral-Mini-3B-2507
|
||||
#
|
||||
# Copyright 2025 Mistral AI and the HuggingFace Inc. team.
|
||||
# Licensed under the Apache License, Version 2.0.
|
||||
"""Inference-only Voxtral (speech-to-text) model."""
|
||||
|
||||
import math
|
||||
from typing import Any, Iterable, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from sglang.srt.layers.activation import get_act_fn
|
||||
from sglang.srt.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.managers.mm_utils import (
|
||||
MultiModalityDataPaddingPatternMultimodalTokens,
|
||||
general_mm_embed_routine,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalInputs,
|
||||
)
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.models.llama import LlamaForCausalLM
|
||||
|
||||
|
||||
class AudioLanguageAdapter(nn.Module):
|
||||
"""MLP projector: Linear -> GELU -> Linear (no bias)."""
|
||||
|
||||
def __init__(self, hidden_size: int, dim: int) -> None:
|
||||
super().__init__()
|
||||
self.w_in = nn.Linear(hidden_size, dim, bias=False)
|
||||
self.gelu = nn.GELU()
|
||||
self.w_out = nn.Linear(dim, dim, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.w_out(self.gelu(self.w_in(x)))
|
||||
|
||||
|
||||
class VoxtralWhisperAttention(nn.Module):
|
||||
"""Multi-headed self-attention using plain SDPA (no KV cache).
|
||||
|
||||
Note: HF Voxtral has bias on q_proj, v_proj, out_proj but NOT on k_proj.
|
||||
We use QKVParallelLinear with bias=True and create a zero bias for k_proj
|
||||
during weight loading.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
num_heads: int,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.head_dim = embed_dim // num_heads
|
||||
self.scaling = self.head_dim**-0.5
|
||||
|
||||
self.qkv_proj = QKVParallelLinear(
|
||||
embed_dim, self.head_dim, num_heads, quant_config=quant_config
|
||||
)
|
||||
# After TP split, the local head count lives on the linear layer
|
||||
self.num_heads = self.qkv_proj.num_heads
|
||||
self.out_proj = RowParallelLinear(
|
||||
embed_dim, embed_dim, bias=True, quant_config=quant_config
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
batch_size, seq_len, _ = hidden_states.shape
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
q, k, v = qkv.chunk(3, dim=-1)
|
||||
q = q * self.scaling
|
||||
|
||||
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(
|
||||
0, 2, 1, 3
|
||||
)
|
||||
|
||||
attn_output = torch.nn.functional.scaled_dot_product_attention(
|
||||
q, k, v, scale=1.0
|
||||
)
|
||||
attn_output = attn_output.permute(0, 2, 1, 3).reshape(
|
||||
batch_size, seq_len, self.num_heads * self.head_dim
|
||||
)
|
||||
attn_output, _ = self.out_proj(attn_output)
|
||||
return attn_output
|
||||
|
||||
|
||||
class VoxtralWhisperEncoderLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
):
|
||||
super().__init__()
|
||||
embed_dim = config.d_model
|
||||
self.self_attn = VoxtralWhisperAttention(
|
||||
embed_dim=embed_dim,
|
||||
num_heads=config.encoder_attention_heads,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
self.self_attn_layer_norm = nn.LayerNorm(embed_dim)
|
||||
self.activation_fn = get_act_fn(
|
||||
getattr(config, "activation_function", "gelu"),
|
||||
quant_config=quant_config,
|
||||
)
|
||||
self.fc1 = ColumnParallelLinear(embed_dim, config.encoder_ffn_dim)
|
||||
self.fc2 = RowParallelLinear(config.encoder_ffn_dim, embed_dim)
|
||||
self.final_layer_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||||
hidden_states = self.self_attn(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.final_layer_norm(hidden_states)
|
||||
hidden_states, _ = self.fc1(hidden_states)
|
||||
hidden_states = self.activation_fn(hidden_states)
|
||||
hidden_states, _ = self.fc2(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
if hidden_states.dtype == torch.float16:
|
||||
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
||||
hidden_states = torch.clamp(
|
||||
hidden_states, min=-clamp_value, max=clamp_value
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class VoxtralWhisperEncoder(nn.Module):
|
||||
"""Whisper encoder (Conv1d + positional embed + transformer + layer norm)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
):
|
||||
super().__init__()
|
||||
embed_dim = config.d_model
|
||||
|
||||
self.conv1 = nn.Conv1d(config.num_mel_bins, embed_dim, kernel_size=3, padding=1)
|
||||
self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
|
||||
self.embed_positions = nn.Embedding(config.max_source_positions, embed_dim)
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
VoxtralWhisperEncoderLayer(config, quant_config)
|
||||
for _ in range(config.encoder_layers)
|
||||
]
|
||||
)
|
||||
self.layer_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
def forward(self, input_features: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
input_features: [batch, num_mel_bins, seq_len]
|
||||
Returns:
|
||||
[batch, seq_len // 2, d_model]
|
||||
"""
|
||||
inputs_embeds = torch.nn.functional.gelu(self.conv1(input_features))
|
||||
inputs_embeds = torch.nn.functional.gelu(self.conv2(inputs_embeds))
|
||||
inputs_embeds = inputs_embeds.permute(0, 2, 1)
|
||||
|
||||
seq_len = inputs_embeds.shape[1]
|
||||
position_ids = torch.arange(seq_len, device=inputs_embeds.device)
|
||||
hidden_states = inputs_embeds + self.embed_positions(position_ids)
|
||||
|
||||
for layer in self.layers:
|
||||
hidden_states = layer(hidden_states)
|
||||
|
||||
hidden_states = self.layer_norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class VoxtralForConditionalGeneration(nn.Module):
|
||||
"""Voxtral: Whisper encoder + MLP projector + Llama decoder.
|
||||
|
||||
HF weight prefixes:
|
||||
audio_tower.* -> self.audio_tower (VoxtralWhisperEncoder)
|
||||
multi_modal_projector.* -> self.multi_modal_projector (AudioLanguageAdapter)
|
||||
language_model.* -> self.language_model (LlamaForCausalLM)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
quant_config: Optional[QuantizationConfig] = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
audio_config = config.audio_config
|
||||
text_config = config.text_config
|
||||
|
||||
# Ensure text_config has rope_parameters (transformers v5 compatibility)
|
||||
if not hasattr(text_config, "rope_parameters"):
|
||||
text_config.rope_parameters = {
|
||||
"rope_type": getattr(text_config, "rope_type", "default"),
|
||||
"rope_theta": getattr(text_config, "rope_theta", 10000.0),
|
||||
}
|
||||
if getattr(text_config, "rope_scaling", None):
|
||||
text_config.rope_parameters.update(text_config.rope_scaling)
|
||||
|
||||
# Infer downsample_factor: intermediate_size / hidden_size for HF format
|
||||
self.downsample_factor = getattr(
|
||||
audio_config,
|
||||
"downsample_factor",
|
||||
audio_config.intermediate_size // audio_config.hidden_size,
|
||||
)
|
||||
|
||||
# Encoder (named audio_tower to match HF weight prefix directly)
|
||||
self.audio_tower = VoxtralWhisperEncoder(audio_config, quant_config)
|
||||
|
||||
# Projector: input = d_model * downsample_factor, output = text_hidden_size
|
||||
adapter_input_dim = audio_config.d_model * self.downsample_factor
|
||||
self.multi_modal_projector = AudioLanguageAdapter(
|
||||
hidden_size=adapter_input_dim,
|
||||
dim=text_config.hidden_size,
|
||||
)
|
||||
|
||||
# Language model
|
||||
self.language_model = LlamaForCausalLM(text_config, quant_config=quant_config)
|
||||
|
||||
# Mel filter bank for raw waveform -> mel spectrogram
|
||||
self._init_mel_filters(audio_config)
|
||||
|
||||
self.pattern = MultiModalityDataPaddingPatternMultimodalTokens()
|
||||
|
||||
def _init_mel_filters(self, audio_config: PretrainedConfig):
|
||||
"""Initialize mel filter bank for mel spectrogram computation."""
|
||||
self._window_size = getattr(audio_config, "window_size", 400)
|
||||
self._hop_length = getattr(audio_config, "hop_length", 160)
|
||||
self._sampling_rate = getattr(audio_config, "sampling_rate", 16000)
|
||||
|
||||
try:
|
||||
from mistral_common.audio import mel_filter_bank
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"mistral_common is required for Voxtral. "
|
||||
"Install it with: pip install mistral_common"
|
||||
)
|
||||
|
||||
mel_filters = mel_filter_bank(
|
||||
num_frequency_bins=1 + self._window_size // 2,
|
||||
num_mel_bins=audio_config.num_mel_bins,
|
||||
min_frequency=0.0,
|
||||
max_frequency=8000.0,
|
||||
sampling_rate=self._sampling_rate,
|
||||
)
|
||||
self.register_buffer(
|
||||
"mel_filters", torch.tensor(mel_filters, dtype=torch.float32)
|
||||
)
|
||||
|
||||
@property
|
||||
def _conv_downsample_factor(self) -> int:
|
||||
return self.audio_tower.conv1.stride[0] * self.audio_tower.conv2.stride[0]
|
||||
|
||||
@property
|
||||
def _chunk_size(self) -> int:
|
||||
return (
|
||||
self.config.audio_config.max_source_positions * self._conv_downsample_factor
|
||||
)
|
||||
|
||||
def _compute_mel_spectrogram(self, audio_waveform: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute log-mel spectrogram from raw waveform using STFT."""
|
||||
window = torch.hann_window(self._window_size, device=audio_waveform.device)
|
||||
stft = torch.stft(
|
||||
audio_waveform,
|
||||
self._window_size,
|
||||
self._hop_length,
|
||||
window=window,
|
||||
return_complex=True,
|
||||
)
|
||||
magnitudes = stft[..., :-1].abs() ** 2
|
||||
mel_spec = self.mel_filters.T @ magnitudes
|
||||
log_spec = torch.clamp(mel_spec, min=1e-10).log10()
|
||||
log_spec_max = log_spec.max()
|
||||
log_spec = torch.maximum(log_spec, log_spec_max - 8.0)
|
||||
log_spec = (log_spec + 4.0) / 4.0
|
||||
return log_spec
|
||||
|
||||
def _encode_audio(self, audio_waveforms: List[torch.Tensor]) -> List[torch.Tensor]:
|
||||
"""Encode raw audio waveforms through mel spectrogram + whisper encoder."""
|
||||
dtype = self.audio_tower.conv1.weight.dtype
|
||||
device = self.audio_tower.conv1.weight.device
|
||||
|
||||
chunked_features: List[torch.Tensor] = []
|
||||
chunks_per_example: List[int] = []
|
||||
chunk_size = self._chunk_size
|
||||
# Pad raw audio to a multiple of chunk_samples so that silence is
|
||||
# properly converted to mel features (matching HF VoxtralProcessor).
|
||||
chunk_samples = chunk_size * self._hop_length
|
||||
|
||||
for waveform in audio_waveforms:
|
||||
waveform = waveform.to(device=device, dtype=torch.float32)
|
||||
n_samples = waveform.shape[-1]
|
||||
target_samples = chunk_samples * math.ceil(n_samples / chunk_samples)
|
||||
if target_samples > n_samples:
|
||||
waveform = torch.nn.functional.pad(
|
||||
waveform, (0, target_samples - n_samples)
|
||||
)
|
||||
mel = self._compute_mel_spectrogram(waveform)
|
||||
chunks = mel.split(chunk_size, dim=-1)
|
||||
chunked_features.extend(chunks)
|
||||
chunks_per_example.append(len(chunks))
|
||||
|
||||
if not chunked_features:
|
||||
return []
|
||||
|
||||
input_embeds = torch.stack(chunked_features).to(dtype)
|
||||
encoder_out = self.audio_tower(input_embeds)
|
||||
|
||||
results = []
|
||||
chunk_idx = 0
|
||||
for n_chunks in chunks_per_example:
|
||||
result = encoder_out[chunk_idx : chunk_idx + n_chunks].flatten(0, 1)
|
||||
results.append(result)
|
||||
chunk_idx += n_chunks
|
||||
|
||||
return results
|
||||
|
||||
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
|
||||
return self.pattern.pad_input_tokens(input_ids, mm_inputs)
|
||||
|
||||
def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||
"""Encode audio waveforms -> downsample -> project."""
|
||||
audio_waveforms = [item.feature for item in items]
|
||||
audio_embeddings = self._encode_audio(audio_waveforms)
|
||||
|
||||
# Downsample: reshape to merge adjacent frames
|
||||
for i, emb in enumerate(audio_embeddings):
|
||||
seq_len, dim = emb.shape
|
||||
audio_embeddings[i] = emb.reshape(
|
||||
seq_len // self.downsample_factor,
|
||||
dim * self.downsample_factor,
|
||||
)
|
||||
|
||||
# Project through adapter
|
||||
packed = torch.cat(audio_embeddings, dim=0)
|
||||
packed = self.multi_modal_projector(packed)
|
||||
|
||||
return packed
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
hidden_states = general_mm_embed_routine(
|
||||
input_ids=input_ids,
|
||||
forward_batch=forward_batch,
|
||||
language_model=self.language_model,
|
||||
data_embedding_funcs={
|
||||
Modality.AUDIO: self.get_audio_feature,
|
||||
},
|
||||
positions=positions,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
def get_language_model(self) -> nn.Module:
|
||||
return self.language_model
|
||||
|
||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||
encoder_stacked = [
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
("qkv_proj", "v_proj", "v"),
|
||||
]
|
||||
|
||||
encoder_dict = dict(self.audio_tower.named_parameters())
|
||||
projector_dict = dict(self.multi_modal_projector.named_parameters())
|
||||
|
||||
# Collect all weights; synthesise missing k_proj bias as zeros.
|
||||
weights_list = list(weights)
|
||||
extra_weights = []
|
||||
for name, w in weights_list:
|
||||
if name.startswith("audio_tower.") and ".self_attn.k_proj.weight" in name:
|
||||
bias_name = name.replace(".weight", ".bias")
|
||||
if not any(n == bias_name for n, _ in weights_list):
|
||||
extra_weights.append(
|
||||
(bias_name, torch.zeros(w.shape[0], dtype=w.dtype))
|
||||
)
|
||||
weights_list.extend(extra_weights)
|
||||
|
||||
def llm_weights_generator():
|
||||
for name, w in weights_list:
|
||||
# Encoder weights
|
||||
if name.startswith("audio_tower."):
|
||||
trimmed = name[len("audio_tower.") :]
|
||||
loaded = False
|
||||
for param_name, weight_name, shard_id in encoder_stacked:
|
||||
if f".{weight_name}." in trimmed:
|
||||
stacked_name = trimmed.replace(weight_name, param_name)
|
||||
if stacked_name in encoder_dict:
|
||||
param = encoder_dict[stacked_name]
|
||||
param.weight_loader(param, w, shard_id)
|
||||
loaded = True
|
||||
break
|
||||
if not loaded and trimmed in encoder_dict:
|
||||
param = encoder_dict[trimmed]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", default_weight_loader
|
||||
)
|
||||
weight_loader(param, w)
|
||||
continue
|
||||
|
||||
# Projector weights
|
||||
if name.startswith("multi_modal_projector."):
|
||||
trimmed = name[len("multi_modal_projector.") :]
|
||||
trimmed = trimmed.replace("linear_1.", "w_in.").replace(
|
||||
"linear_2.", "w_out."
|
||||
)
|
||||
if trimmed in projector_dict:
|
||||
param = projector_dict[trimmed]
|
||||
default_weight_loader(param, w)
|
||||
continue
|
||||
|
||||
# LLM weights
|
||||
if name.startswith("language_model."):
|
||||
name = name[len("language_model.") :]
|
||||
yield (name, w)
|
||||
|
||||
self.language_model.load_weights(llm_weights_generator())
|
||||
|
||||
|
||||
EntryClass = [VoxtralForConditionalGeneration]
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Multimodal processor for Voxtral (speech-to-text) models."""
|
||||
|
||||
import math
|
||||
import re
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
MultimodalDataItem,
|
||||
MultimodalProcessorOutput,
|
||||
)
|
||||
from sglang.srt.models.voxtral import VoxtralForConditionalGeneration
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
BaseMultimodalProcessor,
|
||||
MultimodalSpecialTokens,
|
||||
)
|
||||
|
||||
# Special token IDs for Voxtral audio (from tekken.json vocabulary)
|
||||
AUDIO_TOKEN_ID = 24 # [AUDIO]
|
||||
BEGIN_AUDIO_TOKEN_ID = 25 # [BEGIN_AUDIO]
|
||||
INST_TOKEN_ID = 3 # [INST]
|
||||
|
||||
# Placeholder for load_mm_data regex matching.
|
||||
# encode("[AUDIO]") does NOT produce token 24; actual token insertion
|
||||
# is handled in _build_input_ids_with_audio.
|
||||
AUDIO_PLACEHOLDER = "[AUDIO]"
|
||||
AUDIO_PLACEHOLDER_REGEX = re.compile(r"\[AUDIO\]")
|
||||
|
||||
|
||||
class VoxtralMultimodalProcessor(BaseMultimodalProcessor):
|
||||
models = [VoxtralForConditionalGeneration]
|
||||
|
||||
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
|
||||
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||
audio_config = getattr(hf_config, "audio_config", None)
|
||||
self.audio_token_id = getattr(hf_config, "audio_token_id", AUDIO_TOKEN_ID)
|
||||
self.sampling_rate = getattr(audio_config, "sampling_rate", 16000)
|
||||
self.hop_length = getattr(audio_config, "hop_length", 160)
|
||||
self.max_source_positions = getattr(audio_config, "max_source_positions", 1500)
|
||||
self.conv_downsample = 2 # conv1 stride=1 * conv2 stride=2
|
||||
self.downsample_factor = getattr(
|
||||
audio_config,
|
||||
"downsample_factor",
|
||||
getattr(audio_config, "intermediate_size", 5120)
|
||||
// getattr(audio_config, "hidden_size", 1280),
|
||||
)
|
||||
|
||||
self.mm_tokens = MultimodalSpecialTokens(
|
||||
audio_token=AUDIO_PLACEHOLDER,
|
||||
audio_token_regex=AUDIO_PLACEHOLDER_REGEX,
|
||||
audio_token_id=self.audio_token_id,
|
||||
).build(_processor)
|
||||
|
||||
def _compute_audio_token_count(self, n_samples: int) -> int:
|
||||
"""Compute the number of [AUDIO] tokens for a given audio length."""
|
||||
mel_frames = n_samples / self.hop_length
|
||||
chunk_size = self.max_source_positions * self.conv_downsample
|
||||
n_chunks = math.ceil(mel_frames / chunk_size) if mel_frames > 0 else 1
|
||||
tokens_per_chunk = self.max_source_positions // self.downsample_factor
|
||||
return n_chunks * tokens_per_chunk
|
||||
|
||||
async def process_mm_data_async(
|
||||
self,
|
||||
image_data,
|
||||
audio_data,
|
||||
input_text,
|
||||
request_obj,
|
||||
**kwargs,
|
||||
) -> Optional[MultimodalProcessorOutput]:
|
||||
if not audio_data:
|
||||
return None
|
||||
|
||||
# Insert [AUDIO] placeholders into prompt for load_mm_data's regex
|
||||
prompt_with_placeholders = self._insert_audio_placeholders(
|
||||
input_text, len(audio_data)
|
||||
)
|
||||
|
||||
# load_mm_data handles async loading, format detection, resampling.
|
||||
# process_and_combine_mm_data cannot be used: HF VoxtralProcessor.__call__
|
||||
# does not support audio (only apply_chat_template does).
|
||||
base_output = self.load_mm_data(
|
||||
prompt=prompt_with_placeholders,
|
||||
audio_data=audio_data,
|
||||
multimodal_tokens=self.mm_tokens,
|
||||
audio_sample_rate=self.sampling_rate,
|
||||
)
|
||||
if base_output is None:
|
||||
return None
|
||||
|
||||
# Convert loaded audio to tensors
|
||||
waveforms: List[torch.Tensor] = []
|
||||
for audio in base_output.audios:
|
||||
wav = torch.as_tensor(audio, dtype=torch.float32)
|
||||
if wav.dim() > 1:
|
||||
wav = wav.mean(dim=0)
|
||||
waveforms.append(wav)
|
||||
|
||||
# Compute audio token counts and build input_ids with audio tokens
|
||||
audio_token_counts = [
|
||||
self._compute_audio_token_count(wav.shape[-1]) for wav in waveforms
|
||||
]
|
||||
tokenizer = getattr(self._processor, "tokenizer", self._processor)
|
||||
input_ids = self._build_input_ids_with_audio(
|
||||
tokenizer, input_text, audio_token_counts
|
||||
)
|
||||
|
||||
# Find offsets of [AUDIO] token runs and build mm_items
|
||||
audio_offsets = self._find_audio_offsets(input_ids, self.audio_token_id)
|
||||
mm_items = []
|
||||
for i, wav in enumerate(waveforms):
|
||||
item = MultimodalDataItem(feature=wav, modality=Modality.AUDIO)
|
||||
if i < len(audio_offsets):
|
||||
item.offsets = [audio_offsets[i]]
|
||||
mm_items.append(item)
|
||||
|
||||
return MultimodalProcessorOutput(
|
||||
input_ids=input_ids,
|
||||
mm_items=mm_items,
|
||||
audio_token_id=self.audio_token_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _insert_audio_placeholders(prompt: str, n_audio: int) -> str:
|
||||
"""Insert [AUDIO] placeholder texts into the prompt for load_mm_data."""
|
||||
placeholders = AUDIO_PLACEHOLDER * n_audio
|
||||
# Insert after the last [INST] marker if present
|
||||
last_inst = prompt.rfind("[INST]")
|
||||
if last_inst >= 0:
|
||||
insert_pos = last_inst + len("[INST]")
|
||||
return prompt[:insert_pos] + placeholders + prompt[insert_pos:]
|
||||
return placeholders + prompt
|
||||
|
||||
@staticmethod
|
||||
def _find_audio_offsets(input_ids: List[int], audio_token_id: int) -> List[tuple]:
|
||||
"""Find consecutive runs of audio_token_id in input_ids."""
|
||||
offsets = []
|
||||
start = None
|
||||
for i, tok_id in enumerate(input_ids):
|
||||
if tok_id == audio_token_id:
|
||||
if start is None:
|
||||
start = i
|
||||
elif start is not None:
|
||||
offsets.append((start, i - 1))
|
||||
start = None
|
||||
if start is not None:
|
||||
offsets.append((start, len(input_ids) - 1))
|
||||
return offsets
|
||||
|
||||
def _build_input_ids_with_audio(
|
||||
self,
|
||||
tokenizer,
|
||||
input_text: str,
|
||||
audio_token_counts: List[int],
|
||||
) -> List[int]:
|
||||
"""Build input_ids by tokenizing text and inserting audio tokens.
|
||||
|
||||
The input_text is a decoded Mistral prompt (from text-only
|
||||
apply_chat_template). We re-tokenize to get proper special tokens
|
||||
(BOS, [INST], [/INST]), then insert [BEGIN_AUDIO] + [AUDIO]*N after
|
||||
the last [INST].
|
||||
"""
|
||||
messages = self._parse_mistral_prompt(input_text)
|
||||
try:
|
||||
input_ids = tokenizer.apply_chat_template(messages, tokenize=True)
|
||||
except (ValueError, KeyError):
|
||||
# Fallback if prompt parsing produces malformed messages
|
||||
input_ids = tokenizer.encode(input_text)
|
||||
|
||||
# Insert audio tokens after the last [INST]
|
||||
inst_positions = [i for i, t in enumerate(input_ids) if t == INST_TOKEN_ID]
|
||||
insert_pos = (inst_positions[-1] + 1) if inst_positions else 1
|
||||
|
||||
audio_tokens = []
|
||||
for count in audio_token_counts:
|
||||
audio_tokens.append(BEGIN_AUDIO_TOKEN_ID)
|
||||
audio_tokens.extend([AUDIO_TOKEN_ID] * count)
|
||||
|
||||
return input_ids[:insert_pos] + audio_tokens + input_ids[insert_pos:]
|
||||
|
||||
@staticmethod
|
||||
def _parse_mistral_prompt(prompt: str) -> List[Dict[str, str]]:
|
||||
"""Parse a Mistral-formatted prompt into a list of messages."""
|
||||
messages = []
|
||||
text = prompt.strip()
|
||||
|
||||
for marker in ["<s>", "</s>"]:
|
||||
text = text.replace(marker, "")
|
||||
text = text.strip()
|
||||
|
||||
# Extract system prompt
|
||||
system_match = re.search(
|
||||
r"\[SYSTEM_PROMPT\]\s*(.*?)\s*\[/SYSTEM_PROMPT\]", text, re.DOTALL
|
||||
)
|
||||
if system_match:
|
||||
messages.append(
|
||||
{"role": "system", "content": system_match.group(1).strip()}
|
||||
)
|
||||
text = text[: system_match.start()] + text[system_match.end() :]
|
||||
text = text.strip()
|
||||
|
||||
# Split by [INST] / [/INST]
|
||||
parts = re.split(r"\[/?INST\]", text)
|
||||
for i, part in enumerate(parts):
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
if i % 2 == 1:
|
||||
messages.append({"role": "user", "content": part})
|
||||
elif i > 0:
|
||||
messages.append({"role": "assistant", "content": part})
|
||||
|
||||
if not messages:
|
||||
messages.append({"role": "user", "content": text})
|
||||
|
||||
return messages
|
||||
@@ -837,9 +837,25 @@ def get_tokenizer(
|
||||
)
|
||||
raise RuntimeError(err_msg) from e
|
||||
except ValueError as e:
|
||||
# MistralCommon tokenizers reject standard HF kwargs like
|
||||
# trust_remote_code, use_fast etc. Retry without them.
|
||||
if "are not supported by" in str(e) and "MistralCommon" in str(e):
|
||||
for k in (
|
||||
"trust_remote_code",
|
||||
"tokenizer_revision",
|
||||
"use_fast",
|
||||
"_from_auto",
|
||||
"clean_up_tokenization_spaces",
|
||||
):
|
||||
kwargs.pop(k, None)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_name,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
# If the error pertains to the tokenizer class not existing or not
|
||||
# currently being imported, suggest using the --trust-remote-code flag.
|
||||
if not trust_remote_code and (
|
||||
elif not trust_remote_code and (
|
||||
"does not exist or is not currently imported." in str(e)
|
||||
or "requires you to execute the tokenizer file" in str(e)
|
||||
):
|
||||
@@ -875,6 +891,7 @@ def get_tokenizer(
|
||||
"slowdown. Consider using a fast tokenizer instead."
|
||||
)
|
||||
|
||||
_patch_mistral_common_tokenizer(tokenizer)
|
||||
_fix_special_tokens_pattern(tokenizer)
|
||||
attach_additional_stop_token_ids(tokenizer)
|
||||
tokenizer = patch_tokenizer(tokenizer)
|
||||
@@ -1267,6 +1284,22 @@ def get_processor(
|
||||
revision=revision,
|
||||
**kwargs,
|
||||
)
|
||||
elif (
|
||||
"are not supported by" in error_message and "MistralCommon" in error_message
|
||||
):
|
||||
logger.info(
|
||||
"AutoProcessor for %s rejected standard kwargs, "
|
||||
"retrying without trust_remote_code/use_fast",
|
||||
tokenizer_name,
|
||||
)
|
||||
kwargs.pop("use_fast", None)
|
||||
kwargs.pop("_from_auto", None)
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
tokenizer_name,
|
||||
*args,
|
||||
revision=revision,
|
||||
**kwargs,
|
||||
)
|
||||
elif "Unrecognized feature extractor" in error_message:
|
||||
logger.info(
|
||||
"AutoProcessor failed on feature extractor for %s, "
|
||||
@@ -1315,6 +1348,7 @@ def get_processor(
|
||||
)
|
||||
|
||||
tokenizer = get_tokenizer_from_processor(processor)
|
||||
_patch_mistral_common_tokenizer(tokenizer)
|
||||
|
||||
if tokenizer.chat_template is None:
|
||||
local_path = download_from_hf(
|
||||
@@ -1341,6 +1375,86 @@ def attach_additional_stop_token_ids(tokenizer):
|
||||
tokenizer.additional_stop_token_ids = None
|
||||
|
||||
|
||||
def _patch_mistral_common_tokenizer(tokenizer):
|
||||
"""Patch MistralCommonTokenizer/Backend to be compatible with HF tokenizer API.
|
||||
|
||||
MistralCommon tokenizers (used by Voxtral, Pixtral, etc.) reject several
|
||||
standard kwargs and lack some attributes that sglang expects. We wrap the
|
||||
offending methods once at load time so that the rest of the codebase does
|
||||
not need any special-casing.
|
||||
"""
|
||||
cls_name = type(tokenizer).__name__
|
||||
if "MistralCommon" not in cls_name:
|
||||
return tokenizer
|
||||
if getattr(tokenizer, "_mistral_common_patched", False):
|
||||
return tokenizer
|
||||
tokenizer._mistral_common_patched = True
|
||||
|
||||
# Missing attributes
|
||||
if not hasattr(tokenizer, "get_added_vocab"):
|
||||
tokenizer.get_added_vocab = lambda: {}
|
||||
|
||||
# Set a chat_template containing "audio" so that sglang's content format
|
||||
# detector returns "openai" (which preserves audio_url extraction).
|
||||
# The actual template rendering is done by MistralCommon's apply_chat_template.
|
||||
if not hasattr(tokenizer, "chat_template") or tokenizer.chat_template is None:
|
||||
tokenizer.chat_template = "<!-- audio/image multimodal -->"
|
||||
|
||||
# convert_tokens_to_ids asserts on multi-token strings
|
||||
_orig_convert = tokenizer.convert_tokens_to_ids
|
||||
|
||||
def _safe_convert(val):
|
||||
try:
|
||||
return _orig_convert(val)
|
||||
except AssertionError:
|
||||
return getattr(tokenizer, "unk_token_id", None)
|
||||
|
||||
tokenizer.convert_tokens_to_ids = _safe_convert
|
||||
|
||||
# Wrap methods that reject certain kwargs
|
||||
def _drop_kwargs(fn, keys):
|
||||
def wrapper(*args, **kwargs):
|
||||
for k in keys:
|
||||
kwargs.pop(k, None)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
tokenizer.decode = _drop_kwargs(tokenizer.decode, ["spaces_between_special_tokens"])
|
||||
tokenizer.batch_decode = _drop_kwargs(
|
||||
tokenizer.batch_decode, ["spaces_between_special_tokens"]
|
||||
)
|
||||
|
||||
# Save original apply_chat_template for processors that need it (e.g. Voxtral)
|
||||
tokenizer._orig_apply_chat_template = tokenizer.apply_chat_template
|
||||
|
||||
def _safe_apply_chat_template(messages, **kwargs):
|
||||
"""Wrapper that strips unsupported kwargs and non-text content parts.
|
||||
|
||||
When sglang extracts audio/image URLs, it replaces content blocks with
|
||||
{"type": "audio"} or {"type": "image"} (no URL). MistralCommon fails
|
||||
on these stripped blocks. We convert them to text-only messages.
|
||||
"""
|
||||
kwargs.pop("add_generation_prompt", None)
|
||||
cleaned = []
|
||||
for msg in messages:
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, list):
|
||||
text_parts = [
|
||||
p.get("text", "")
|
||||
for p in content
|
||||
if isinstance(p, dict) and p.get("type") == "text"
|
||||
]
|
||||
msg = {**msg, "content": " ".join(text_parts) if text_parts else ""}
|
||||
cleaned.append(msg)
|
||||
else:
|
||||
cleaned.append(msg)
|
||||
return tokenizer._orig_apply_chat_template(cleaned, **kwargs)
|
||||
|
||||
tokenizer.apply_chat_template = _safe_apply_chat_template
|
||||
|
||||
|
||||
def check_gguf_file(model: Union[str, os.PathLike]) -> bool:
|
||||
"""Check if the file is a GGUF model."""
|
||||
model = Path(model)
|
||||
|
||||
Reference in New Issue
Block a user