vllm-cpu 0.11.0.post2__cp312-cp312-manylinux_2_17_x86_64.whl

vllm/model_executor/models/voxtral.py

@@ -0,0 +1,786 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from collections.abc import Iterable, Mapping, Sequence
+from functools import cached_property
+from math import ceil
+from typing import Literal, Optional, Union, cast
+
+import numpy as np
+import regex as re
+import torch
+import torch.nn as nn
+from mistral_common.audio import mel_filter_bank
+from mistral_common.protocol.instruct.messages import (AudioChunk, RawAudio,
+                                                       TextChunk, UserMessage)
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.protocol.transcription.request import TranscriptionRequest
+from mistral_common.tokens.tokenizers.audio import Audio, AudioEncoder
+from transformers import BatchFeature, TensorType, WhisperConfig
+from transformers.tokenization_utils_base import TextInput
+
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
+from vllm.inputs.data import PromptType
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models import SupportsPP
+from vllm.model_executor.models.module_mapping import MultiModelKeys
+# yapf: disable
+from vllm.model_executor.models.whisper import WhisperEncoder
+# yapf: enable
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
+                                    MultiModalKwargsItems, MultiModalUUIDDict,
+                                    NestedTensors)
+from vllm.multimodal.parse import (AudioProcessorItems, MultiModalDataItems,
+                                   MultiModalDataParser)
+from vllm.multimodal.processing import (BaseMultiModalProcessor,
+                                        BaseProcessingInfo,
+                                        MultiModalProcessingInfo,
+                                        PromptReplacement, PromptUpdate)
+from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.tokenizer import (MistralTokenizer,
+                                               cached_tokenizer_from_config)
+
+from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
+                         SupportsMultiModal, SupportsTranscription)
+from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,
+                    merge_multimodal_embeddings)
+
+logger = init_logger(__name__)
+
+ISO639_1_SUPPORTED_LANGS = {
+    "ar": "Arabic",
+    "nl": "Dutch",
+    "en": "English",
+    "fr": "French",
+    "de": "German",
+    "hi": "Hindi",
+    "it": "Italian",
+    "pt": "Portuguese",
+    "es": "Spanish",
+}
+
+
+class VoxtralProcessorAdapter:
+    """
+    Provide a HF-compatible interface for
+    :class:`mistral_common.tokens.tokenizers.multimodal.AudioEncoder`.
+    """
+
+    def __init__(self, tokenizer: MistralTokenizer) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+
+    @cached_property
+    def _audio_processor(self) -> AudioEncoder:
+        audio_encoder = self.tokenizer.instruct.audio_encoder
+        assert isinstance(audio_encoder, AudioEncoder)
+        return audio_encoder
+
+    @cached_property
+    def audio_token_id(self) -> int:
+        return self._audio_processor.special_ids.audio
+
+    @cached_property
+    def begin_audio_token_id(self) -> int:
+        return self._audio_processor.special_ids.begin_audio
+
+    # @cached_property
+    # def begin_transcript_token_id(self) -> int:
+    #     return self._audio_processor.special_ids.begin_transcript
+
+    # @cached_property
+    # def end_transcript_token_id(self) -> int:
+    #     return self._audio_processor.special_ids.end_transcript
+
+    @cached_property
+    def sampling_rate(self) -> int:
+        return self._audio_processor.audio_config.sampling_rate
+
+    @cached_property
+    def frame_rate(self) -> float:
+        return self._audio_processor.audio_config.frame_rate
+
+    def get_num_audio_tokens(
+        self,
+        audio_length: int,
+    ) -> int:
+        pad_audio_length = self._audio_processor.next_multiple_of_chunk_frames(
+            audio_length, self.sampling_rate)
+        return ceil(pad_audio_length / (self.sampling_rate // self.frame_rate))
+
+    def __call__(
+        self,
+        text: Optional[Union[TextInput, list[TextInput]]] = None,
+        audios: Optional[Union[np.ndarray, list[np.ndarray]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> Mapping[str, NestedTensors]:
+        if text is None:
+            text = []
+        if not isinstance(text, list):
+            text = [text]
+        if audios is None:
+            audios = []
+        if not isinstance(audios, list):
+            audios = [audios]
+
+        if not audios:
+            input_ids = self.tokenizer(text).input_ids
+            return {"input_ids": torch.tensor(input_ids)}
+
+        # Allow dummy text, which is used for profiling as well as token inputs
+        if any(len(t) > 0 for t in text):
+            raise ValueError(
+                "You've passed text inputs instead of token inputs. "
+                "Make sure to process your input via `mistral_common`'s "
+                "tokenizer or pass a chat completion request. "
+                "For more info, see: "
+                "https://github.com/vllm-project/vllm/issues/8411.")
+
+        audios_tokens = list[torch.Tensor]()
+        audios_processed = list[torch.Tensor]()
+        for audio in audios:
+            assert isinstance(audio, np.ndarray)
+            assert audio.ndim == 1
+
+            # pad if necessary
+            audio = self._audio_processor.pad(audio, self.sampling_rate)
+
+            audio_tokens = [
+                self.begin_audio_token_id
+            ] + [self.audio_token_id] * self.get_num_audio_tokens(len(audio))
+
+            audios_tokens.append(torch.tensor(audio_tokens))
+            audios_processed.append(torch.tensor(audio))
+
+        return BatchFeature({
+            "input_ids":
+            torch.cat(audios_tokens)[None].expand(len(text), -1),
+            "audio_arrays":
+            audios_processed,
+        })
+
+
+class VoxtralProcessingInfo(BaseProcessingInfo):
+
+    def get_tokenizer(self) -> MistralTokenizer:
+        tokenizer = cached_tokenizer_from_config(self.ctx.model_config)
+        if not isinstance(tokenizer, MistralTokenizer):
+            raise ValueError("This model requires `--tokenizer-mode mistral`")
+
+        return tokenizer
+
+    def get_hf_processor(self) -> VoxtralProcessorAdapter:
+        return VoxtralProcessorAdapter(self.get_tokenizer())
+
+    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
+        return {"audio": 5}  # Performance tends to degrade after 5
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        return {"audio": self.get_max_audio_tokens()}
+
+    def get_max_audio_tokens(self) -> int:
+        return self.ctx.model_config.max_model_len
+
+    def get_max_audio_array_len(self) -> int:
+        processor = self.get_hf_processor()
+        return self.get_max_audio_tokens() * int(
+            processor.sampling_rate // processor.frame_rate)
+
+
+class VoxtralDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralProcessingInfo]):
+
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return ""
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio", 0)
+
+        target_length = self.info.get_max_audio_array_len()
+
+        return {
+            "audio":
+            self._get_dummy_audios(length=target_length, num_audios=num_audios)
+        }
+
+    def get_dummy_processor_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> ProcessorInputs:
+        tokenizer = self.info.get_tokenizer()
+
+        dummy_text = self.get_dummy_text(mm_counts)
+        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts)
+        dummy_audios = dummy_mm_data.get("audio", [])
+
+        audio_chunks: list[AudioChunk] = []
+        format = "wav"
+        for audio in dummy_audios:
+            audio_item = Audio(
+                audio_array=audio,
+                sampling_rate=self.info.get_hf_processor().sampling_rate,
+                format=format,
+            )
+            chunk = AudioChunk(input_audio=RawAudio.from_audio(audio_item))
+            audio_chunks.append(chunk)
+
+        request = ChatCompletionRequest(messages=[
+            UserMessage(content=[TextChunk(text=dummy_text), *audio_chunks]),
+        ])
+        res = tokenizer.mistral.encode_chat_completion(request)
+        dummy_tokens = res.tokens
+        # whixtral tokenizer adds padding to the audio
+        # so we need to update the audio arrays
+        dummy_mm_data["audio"] = [a.audio_array for a in res.audios]
+
+        return ProcessorInputs(prompt=dummy_tokens, mm_data=dummy_mm_data)
+
+
+class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]
+                                 ):
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: Mapping[str, NestedTensors],
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return dict(audio_arrays=MultiModalFieldConfig.batched("audio"))
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+        audio_id = processor.audio_token_id
+
+        def get_replacement(item_idx: int):
+            audios = mm_items.get_items("audio", AudioProcessorItems)
+            audio_len = audios.get_audio_length(item_idx)
+
+            nb_audio_tokens = processor.get_num_audio_tokens(audio_len)
+
+            return [audio_id] * nb_audio_tokens
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target="",  # Never match the prompt (see below note)
+                replacement=get_replacement,
+            ),
+        ]
+
+    def _cached_apply_hf_processor(
+        self,
+        prompt: Union[str, list[int]],
+        mm_data_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+        mm_uuids: Optional[MultiModalUUIDDict] = None,
+    ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
+        prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(
+            prompt=prompt,
+            mm_data_items=mm_data_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
+            mm_uuids=mm_uuids,
+        )
+
+        # NOTE: The tokens are already inserted by the chat template
+        return prompt_ids, mm_info, True
+
+    def _get_data_parser(self) -> MultiModalDataParser:
+        sampling_rate = self.info.get_hf_processor().sampling_rate
+        return MultiModalDataParser(target_sr=sampling_rate)
+
+
+@MULTIMODAL_REGISTRY.register_processor(VoxtralMultiModalProcessor,
+                                        info=VoxtralProcessingInfo,
+                                        dummy_inputs=VoxtralDummyInputsBuilder)
+class VoxtralForConditionalGeneration(nn.Module, SupportsMultiModal,
+                                      SupportsPP, SupportsLoRA,
+                                      SupportsTranscription):
+    supported_languages = ISO639_1_SUPPORTED_LANGS
+
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"]
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
+
+        # update quant config to so that ignored module and target module names
+        # match the vLLM model names
+        if hasattr(vllm_config, "quant_config"):
+            vllm_config.quant_config = self.maybe_update_quant_config(
+                vllm_config.quant_config)
+
+        config = vllm_config.model_config.hf_config
+        self.config = config
+        self.downsample_factor = self.config.audio_config.downsample_factor
+
+        self.language_model = init_vllm_registered_model(
+            vllm_config=vllm_config,
+            hf_config=config.text_config,
+            prefix=maybe_prefix(prefix, "language_model"),
+        )
+        self.whisper_encoder = VoxtralEncoderModel(
+            vllm_config.with_hf_config(config.audio_config),
+            prefix=maybe_prefix(prefix, "whisper_encoder"),
+        )
+        self.audio_language_adapter = AudioLanguageAdapter(
+            hidden_size=config.audio_config.d_model * self.downsample_factor,
+            dim=config.text_config.hidden_size,
+        )
+
+    def get_language_model(self) -> torch.nn.Module:
+        return self.language_model
+
+    def get_mm_mapping(self) -> MultiModelKeys:
+        """Get module prefix for multimodal models to filter LoRA modules."""
+        return MultiModelKeys.from_string_field(
+            language_model="language_model",
+            connector="audio_language_adapter",
+            tower_model=["whisper_encoder"],
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs: object,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        # NOTE: In v1, inputs_embeds is always generated at model runner, this
+        # condition is for v0 compatibility.
+        elif inputs_embeds is None:
+            audio_embeddings = self.get_multimodal_embeddings(**kwargs)
+            inputs_embeds = self.get_input_embeddings(input_ids,
+                                                      audio_embeddings)
+            input_ids = None
+
+        hidden_states = self.language_model.model(input_ids,
+                                                  positions,
+                                                  intermediate_tensors,
+                                                  inputs_embeds=inputs_embeds)
+
+        return hidden_states
+
+    def get_multimodal_embeddings(
+        self, **kwargs
+    ) -> Union[list[torch.Tensor], torch.Tensor, tuple[torch.Tensor, ...],
+               None]:
+        audio_inputs = self._parse_and_validate_audio_arrays(**kwargs)
+        if audio_inputs is None:
+            return None
+
+        audio_embeddings = self.whisper_encoder(audio_inputs)
+
+        for i, audio_embedding in enumerate(audio_embeddings):
+            seq_len, dim = audio_embedding.shape
+            # Pad such that seq_len is divisible by downsample_factor
+            target_seq_len = self.downsample_factor * math.ceil(
+                seq_len / self.downsample_factor)
+            audio_embedding = torch.nn.functional.pad(
+                audio_embedding,
+                (0, 0, 0, target_seq_len - seq_len),
+            )
+            audio_embeddings[i] = audio_embedding.reshape(
+                target_seq_len // self.downsample_factor,
+                dim * self.downsample_factor)
+
+        # Concat, project and resplit
+        audio_embeddings_packed = torch.cat(audio_embeddings, dim=0)
+        audio_embeddings_packed = self.audio_language_adapter(
+            audio_embeddings_packed)
+        audio_embeddings = torch.split(audio_embeddings_packed,
+                                       [a.shape[0] for a in audio_embeddings],
+                                       dim=0)
+
+        return audio_embeddings
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
+    ) -> torch.Tensor:
+        audio_encoder = self.tokenizer.instruct.audio_encoder
+        audio_tok_id = audio_encoder.audio_token
+
+        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
+        if multimodal_embeddings is not None:
+            inputs_embeds = merge_multimodal_embeddings(
+                input_ids, inputs_embeds, multimodal_embeddings, audio_tok_id)
+        return inputs_embeds
+
+    def _parse_and_validate_audio_arrays(
+            self, **kwargs: object) -> Union[list[torch.Tensor], None]:
+        audio_arrays = kwargs.pop("audio_arrays", None)
+        if audio_arrays is None:
+            return None
+
+        if not isinstance(audio_arrays, (torch.Tensor, list)):
+            raise ValueError("Incorrect type of audio_arrays. "
+                             f"Got type: {type(audio_arrays)}")
+
+        audio_arrays = flatten_bn(audio_arrays)
+        if isinstance(audio_arrays, torch.Tensor):
+            audio_arrays = list(audio_arrays.unbind(0))
+        return audio_arrays
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        return self.language_model.compute_logits(hidden_states)
+
+    @classmethod
+    def get_speech_to_text_config(cls, model_config: ModelConfig,
+                                  task_type: str) -> SpeechToTextConfig:
+        tokenizer = cached_tokenizer_from_config(model_config)
+        audio_config = tokenizer.instruct.audio_encoder.audio_config
+        max_audio_clip_s = audio_config.chunk_length_s
+        sample_rate = audio_config.sampling_rate
+        return SpeechToTextConfig(
+            max_audio_clip_s=max_audio_clip_s,
+            sample_rate=sample_rate,
+            # mistral_common and whisper encoder take care of chunking
+            min_energy_split_window_size=None,
+        )
+
+    @classmethod
+    # for speech-to-text transcription
+    def get_generation_prompt(cls, audio: np.ndarray,
+                              model_config: ModelConfig,
+                              stt_config: SpeechToTextConfig,
+                              language: Optional[str],
+                              task_type: Literal["transcribe", "translate"],
+                              request_prompt: str,
+                              to_language: Optional[str]) -> PromptType:
+        tokenizer = cached_tokenizer_from_config(model_config)
+        audio = Audio(audio, int(stt_config.sample_rate),
+                      format="wav")  # lossless
+        req = TranscriptionRequest(model=model_config.model,
+                                   audio=RawAudio.from_audio(audio),
+                                   language=language)
+
+        tokenized = tokenizer.instruct.encode_transcription(req)
+        audio = (tokenized.audios[0].audio_array, stt_config.sample_rate)
+        prompts_dict = {"multi_modal_data": {"audio": audio}}
+        prompts_dict["prompt_token_ids"] = tokenized.tokens
+        return cast(PromptType, prompts_dict)
+
+    @classmethod
+    def get_num_audio_tokens(cls, audio_duration_s: float,
+                             stt_config: SpeechToTextConfig,
+                             model_config: ModelConfig) -> Optional[int]:
+        """
+        Map from audio duration to number of audio tokens produced by the ASR 
+        model, without running a forward pass.
+        This is used for estimating the amount of processing for this audio.
+        """
+        tokenizer = cached_tokenizer_from_config(model_config)
+        adapter = VoxtralProcessorAdapter(tokenizer)
+        return adapter.get_num_audio_tokens(
+            int(audio_duration_s * stt_config.sample_rate))
+
+    def load_weights(self, weights: Iterable[tuple[str,
+                                                   torch.Tensor]]) -> set[str]:
+        # fmt: off
+        remapping_rules = [
+            (r"mm_whisper_embeddings\.(.*)", r"\1"),
+            (r"audio_language_projection\.(.*)", r"audio_language_adapter.\1"),
+            (r"audio_language_adapter\.0\.weight", r"audio_language_adapter.w_in.weight"),  # noqa: E501
+            (r"audio_language_adapter\.2\.weight", r"audio_language_adapter.w_out.weight"),  # noqa: E501
+        ]
+        # fmt: on
+
+        audio_params = dict(
+            nn.ModuleDict({
+                "audio_language_adapter":
+                self.audio_language_adapter,
+            }).named_parameters())
+
+        loaded_weights = set()
+
+        def llm_weights_generator():
+            nonlocal loaded_weights
+            for name, w in weights:
+                is_encoder = (
+                    name.startswith("mm_whisper_embeddings") and
+                    not name.startswith("mm_whisper_embeddings.tok_embeddings")
+                    and not name.startswith(
+                        "mm_whisper_embeddings.audio_language_projection"))
+
+                for pattern, repl in remapping_rules:
+                    if re.fullmatch(pattern, name):
+                        name = re.sub(pattern, repl, name)
+
+                if is_encoder:
+                    name = self.whisper_encoder.load_weight((name, w))
+                    loaded_weights.add(f"whisper_encoder.{name}")
+                    continue
+
+                if name in audio_params:
+                    param = audio_params[name]
+                    with torch.no_grad():
+                        default_weight_loader(param, w)
+                    loaded_weights.add(name)
+                else:
+                    yield (name, w)
+
+        for name in self.language_model.load_weights(llm_weights_generator()):
+            loaded_weights.add(f"language_model.{name}")
+
+        # potentially manually add position embeddings
+        sin_key = "whisper_encoder.whisper_encoder.embed_positions.weight"
+        if sin_key not in loaded_weights:
+            # make sure we don't hit an error here
+            loaded_weights.add(sin_key)
+
+        return loaded_weights
+
+    def maybe_update_quant_config(
+            self, quant_config: QuantizationConfig) -> QuantizationConfig:
+        """
+        Update quant config to so that ignored module and target module names
+        match the vLLM model names.
+        Right now this is specific for compressed-tensors format and
+        load_format mistral.
+        """
+        remapping_rules = [
+            (r"output", r"language_model.lm_head"),
+            (r"layers\.(\d+)\.attention\.wo",
+             r"language_model.model.layers.\1.self_attn.out_proj"),
+            (r"layers\.(\d+)\.attention\.w(.*)",
+             r"language_model.model.layers.\1.self_attn.\2_proj"),
+            (r"layers\.(\d+)\.feed_forward\.w1",
+             r"language_model.model.layers.\1.mlp.gate_proj"),
+            (r"layers\.(\d+)\.feed_forward\.w2",
+             r"language_model.model.layers.\1.mlp.down_proj"),
+            (r"layers\.(\d+)\.feed_forward\.w3",
+             r"language_model.model.layers.\1.mlp.up_proj"),
+            (r"mm_whisper_embeddings\.whisper_encoder\.transformer\.layers\.(\d+)\.attention.w(.*)",
+             r"whisper_encoder.whisper_encoder.layers.\1.layers.self_attn.\2_proj"
+             ),
+            (r"mm_whisper_embeddings\.whisper_encoder\.transformer\.layers\.(\d+)\.attention.wo",
+             r"whisper_encoder.whisper_encoder.layers.\1.layers.self_attn.out_proj"
+             ),
+            (r"mm_whisper_embeddings\.whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward.w(\d+)",
+             r"whisper_encoder.whisper_encoder.layers.\1.layers.mlp.fc\2"),
+            (r"mm_whisper_embeddings\.whisper_encoder\.conv_layers\.0",
+             r"whisper_encoder.whisper_encoder.conv1"),
+            (r"mm_whisper_embeddings\.whisper_encoder\.conv_layers\.1",
+             r"whisper_encoder.whisper_encoder.conv2"),
+            (r"mm_whisper_embeddings\.audio_language_projection\.0",
+             r"audio_language_adapter.w_in"),
+            (r"mm_whisper_embeddings\.audio_language_projection\.2",
+             r"audio_language_adapter.w_out"),
+        ]
+
+        # Update ignore list
+        if hasattr(quant_config, "ignore"):
+            mistral_ignore = []
+            for name in quant_config.ignore:
+                mistral_name = name
+                for pattern, repl in remapping_rules:
+                    if re.fullmatch(pattern, name):
+                        mistral_name = re.sub(pattern, repl, name)
+                mistral_ignore.append(mistral_name)
+            quant_config.ignore = mistral_ignore
+
+        # Update target list
+        if hasattr(quant_config, "config_groups"):
+            config_groups = quant_config.config_groups
+            for group_name in config_groups:
+                if "targets" in config_groups[group_name]:
+                    targets = []
+                    for name in config_groups[group_name]["targets"]:
+                        mistral_name = name
+                        for pattern, repl in remapping_rules:
+                            if re.fullmatch(pattern, name):
+                                mistral_name = re.sub(pattern, repl, name)
+                        targets.append(mistral_name)
+                config_groups[group_name]["targets"] = targets
+            quant_config.config_groups = config_groups
+
+        return quant_config
+
+
+class AudioLanguageAdapter(nn.Module):
+
+    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 VoxtralEncoderModel(nn.Module):
+    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+
+    # fmt: off
+    mistral_remapping = [
+        (r"whisper_encoder\.conv_layers\.0\.(weight|bias)", r"whisper_encoder.conv1.\1"), # noqa: E501
+        (r"whisper_encoder\.conv_layers\.1\.(weight|bias)", r"whisper_encoder.conv2.\1"), # noqa: E501
+        (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.w([qkv])\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn.\2_proj.\3"), # noqa: E501
+        (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.wo\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn.out_proj.\2"), # noqa: E501
+        (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention_norm\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn_layer_norm.\2"), # noqa: E501
+        (r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w1\.(weight|bias)", r"whisper_encoder.layers.\1.mlp.fc1.\2"), # noqa: E501
+        (r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w2\.(weight|bias)", r"whisper_encoder.layers.\1.mlp.fc2.\2"), # noqa: E501
+        (r"whisper_encoder\.transformer\.layers\.(\d+)\.ffn_norm\.(weight|bias)", r"whisper_encoder.layers.\1.final_layer_norm.\2"), # noqa: E501
+        (r"whisper_encoder\.transformer\.norm\.(weight|bias)", r"whisper_encoder.layer_norm.\1"), # noqa: E501
+    ]
+    # fmt: on
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        *,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = cast(WhisperConfig, vllm_config.model_config.hf_config)
+        self.dtype: torch.dtype = vllm_config.model_config.dtype
+        self.whisper_encoder = WhisperEncoder(vllm_config=vllm_config,
+                                              prefix=maybe_prefix(
+                                                  prefix, "whisper_encoder"),
+                                              init_in_fp32=True)
+        mel_filters = mel_filter_bank(
+            num_frequency_bins=1 + self.config.window_size // 2,
+            num_mel_bins=self.config.num_mel_bins,
+            min_frequency=0.0,
+            max_frequency=8000.0,
+            sampling_rate=self.config.sampling_rate,
+        )
+        self.mel_filters = torch.tensor(mel_filters, dtype=torch.float32)
+
+    def compute_whisper_melspec(
+        self,
+        audio_waveforms: torch.Tensor,
+    ) -> torch.Tensor:
+        input_dtype = audio_waveforms.dtype
+        window = torch.hann_window(self.config.window_size).to(
+            audio_waveforms.device)
+        stft = torch.stft(
+            audio_waveforms,
+            self.config.window_size,
+            self.config.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 = torch.maximum(log_spec, log_spec.max() - 8.0)
+        log_spec = (log_spec + 4.0) / 4.0
+        return log_spec.to(input_dtype)
+
+    @property
+    def downsample_factor(self) -> int:
+        return self.whisper_encoder.conv1.stride[
+            0] * self.whisper_encoder.conv2.stride[0]
+
+    @property
+    def chunk_size(self) -> int:
+        return self.config.max_source_positions * self.downsample_factor
+
+    def prepare_inputs_for_conv(
+        self,
+        audio_waveforms: list[torch.Tensor],
+    ) -> tuple[torch.Tensor, list[int]]:
+        assert isinstance(audio_waveforms, list)
+        # list[num_mel_bins, seq_len]
+        input_features = [
+            self.compute_whisper_melspec(audio).to(self.dtype)
+            for audio in audio_waveforms
+        ]
+
+        chunked_features: list[torch.Tensor] = []
+        chunks_per_example: list[int] = []
+        for feature in input_features:
+            chunks = feature.split(self.chunk_size, dim=-1)
+            chunked_features += chunks
+            chunks_per_example.append(len(chunks))
+
+        # [total_num_chunks, num_mel_bins, chunk_size]
+        return torch.stack(chunked_features), chunks_per_example
+
+    def forward(
+        self, input_features: Union[torch.Tensor, list[torch.Tensor]]
+    ) -> list[torch.Tensor]:
+        if not isinstance(input_features, list):
+            input_features = [input_features]
+
+        # Split long inputs into chunks
+        input_embeds, chunks_per_example = (
+            self.prepare_inputs_for_conv(input_features))
+
+        # [total_num_chunks, ceil(chunk_size / downsample_factor), hidden_size]
+        out = self.whisper_encoder([input_embeds])
+
+        # Re-concatenate the chunks
+        chunk_idx = 0
+        results = []
+        for n_chunks in chunks_per_example:
+            result = out[chunk_idx:chunk_idx + n_chunks].flatten(0, 1)
+            results.append(result)
+            chunk_idx += n_chunks
+
+        return results
+
+    def load_weight(self, weight: tuple[str, torch.Tensor]) -> str:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+
+        name, loaded_weight = weight
+        for pattern, repl in self.mistral_remapping:
+            if re.fullmatch(pattern, name):
+                name = re.sub(pattern, repl, name)
+
+        for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            if weight_name not in name:
+                continue
+            name = name.replace(weight_name, param_name)
+
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            weight_loader(param, loaded_weight, shard_id)
+            break
+        else:
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader",
+                                    default_weight_loader)
+            weight_loader(param, loaded_weight)
+
+        return name