Update vLLM link to PyPI and remove bundled vllm_plugin folder (#4)

- Update vLLM link to PyPI and remove bundled vllm_plugin folder (2f6ce47050eddd7054bdc902472eb042fe95fe1f)

README.md

@@ -51,7 +51,7 @@ tags:
 
 <p align="center">
   <a href="https://meralion.org/demo/">💻 Web Demo</a> |
-  <a href="https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/vllm_plugin_meralion2/readme.md">⚙️ vLLM</a>
+  <a href="https://pypi.org/project/vllm-plugin-meralion2/#description">⚙️ vLLM</a>
 </p>
 
 ## Introduction 

vllm_plugin_meralion2/offline_example.py

@@ -1,50 +0,0 @@
-import os
-import torch
-import librosa
-
-from vllm import LLM, SamplingParams
-from vllm_plugin_meralion2 import NoRepeatNGramLogitsProcessor
-
-
-model_name = "MERaLiON/MERaLiON-2-10B"
-# model_name = "MERaLiON/MERaLiON-2-10B-ASR"
-# model_name = "MERaLiON/MERaLiON-2-3B"
-
-llm = LLM(model=model_name,
-          tokenizer=model_name,
-          limit_mm_per_prompt={"audio": 1},
-          trust_remote_code=True,
-          dtype=torch.bfloat16
-          )
-
-# change example.wav to your audio file.
-audio_array, sample_rate = librosa.load("example.wav", sr=16000)
-
-question= "Please trancribe this speech."
-prompt = (
-    "<start_of_turn>user\n"
-    f"Instruction: {question} \nFollow the text instruction based on the following audio: <SpeechHere><end_of_turn>\n"
-    "<start_of_turn>model\n")
-
-sampling_params = SamplingParams(
-  temperature=0.0,
-  top_p=0.9,
-  top_k=50,
-  repetition_penalty=1.0,
-  seed=42,
-  max_tokens=1024,
-  stop_token_ids=None,
-  logits_processors=[NoRepeatNGramLogitsProcessor(6)]
-)
-
-mm_data = {"audio": [(audio_array, sample_rate)]}
-inputs = {"prompt": prompt, "multi_modal_data": mm_data}
-
-# batch inference
-inputs = [inputs] * 2
-
-outputs = llm.generate(inputs, sampling_params=sampling_params)
-
-for o in outputs:
-    generated_text = o.outputs[0].text
-    print(generated_text)

vllm_plugin_meralion2/openai_client_curl.sh

@@ -1,45 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-# target url and port
-TARGET_API=http://localhost:8000
-
-# MERaLiON/MERaLiON-2-10B or MERaLiON/MERaLiON-2-10B-ARS or MERaLiON/MERaLiON-2-3B
-MODEL_NAME=MERaLiON/MERaLiON-2-10B  
-
-# Refer to https://huggingface.co/MERaLiON/MERaLiON-2-10B-ASR#audio-input and https://huggingface.co/MERaLiON/MERaLiON-2-10B#audio-input
-PROMPT='Instruction: Please transcribe this speech. \nFollow the text instruction based on the following audio: <SpeechHere>'
-
-# change to true if need stream output
-STREAM=false
-
-# change example.wav to your audio file.
-base64 -w 0 -i example.wav > audio.b64
-
-
-cat > payload.json <<EOF
-{
-  "model": "${MODEL_NAME}",
-  "messages":[
-    {"role":"user",
-     "content":[
-       {"type":"text","text":"${PROMPT}"},
-       {"type":"audio_url","audio_url":{"url":"data:audio/ogg;base64,$(cat audio.b64)"}}
-     ]
-    }
-  ],
-  "max_completion_tokens":1024,
-  "temperature":0.1,
-  "top_p":0.9,
-  "top_k":50,
-  "repetition_penalty":1.0,
-  "length_penalty":1.0,
-  "logits_processors":[{"qualname":"vllm_plugin_meralion2.NoRepeatNGramLogitsProcessor","args":[6]}],
-  "seed":42,
-  "stream":${STREAM}
-}
-EOF
-
-curl $TARGET_API/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  --data-binary @payload.json

vllm_plugin_meralion2/openai_client_example.py

@@ -1,70 +0,0 @@
-import base64
-
-from openai import OpenAI
-
-prompt_template = "Instruction: {text_input} \nFollow the text instruction based on the following audio: <SpeechHere>"
-
-def get_client(api_key="EMPTY", base_url="http://localhost:8000/v1"):
-    client = OpenAI(
-        api_key=api_key,
-        base_url=base_url,
-    )
-
-    models = client.models.list()
-    model_name = models.data[0].id
-    return client, model_name
-
-
-def get_response(text_input, base64_audio_input=None, **params):
-    if base64_audio_input:
-        content = [
-            {
-                "type": "text",
-                "text": prompt_template.format(text_input=text_input)
-            },
-            {
-                "type": "audio_url",
-                "audio_url": {
-                    "url": f"data:audio/ogg;base64,{base64_audio_input}"
-                },
-            },
-        ]
-    else:
-        content = text_input
-    
-
-    response_obj = client.chat.completions.create(
-        messages=[{
-            "role": "user",
-            "content": content,
-        }],
-        **params
-    )
-    return response_obj
-
-
-# change example.wav to your audio file. Make sure its 16khz, mono channel.
-audio_bytes = open("example.wav", "rb").read()
-audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
-
-# use the port number of your vllm service.
-client, model_name = get_client(base_url="http://localhost:8000/v1")
-
-generation_parameters = dict(
-    model=model_name,
-    max_completion_tokens=1024,
-    temperature=0.0,
-    top_p=0.9,
-    extra_body={
-        "repetition_penalty": 1.0,
-        "top_k": 50,
-        "length_penalty": 1.0,
-        "logits_processors": [
-            {"qualname": "vllm_plugin_meralion2.NoRepeatNGramLogitsProcessor", "args": [6]}
-        ]
-    },
-    seed=42
-)
-
-response_obj = get_response("Please transcribe this speech.", audio_base64, **generation_parameters)
-print(response_obj.choices[0].message.content)

vllm_plugin_meralion2/openai_serve_example.sh

@@ -1,5 +0,0 @@
-vllm serve MERaLiON/MERaLiON-2-10B \
-    --trust-remote-code \
-    --dtype bfloat16 \
-    --logits-processor-pattern vllm_plugin_meralion2.NoRepeatNGramLogitsProcessor \
-    --port 8000

vllm_plugin_meralion2/pyproject.toml

@@ -1,40 +0,0 @@
-# pyproject.toml
-
-[build-system]
-requires = [
-  "setuptools>=61.0",
-  "wheel",
-]
-build-backend = "setuptools.build_meta"
-
-[project]
-name = "vllm_plugin_meralion2"
-version = "0.1.2.post1"
-description = "A vLLM plugin to register the MERaLiON-2-10B model architecture with vLLM’s plugin system."
-authors = [{name = "MERaLiON Team"},]
-readme = "readme.md"
-
-# Python version compatibility
-requires-python = ">=3.9"
-
-# your runtime dependencies
-dependencies = [
-  "vllm>=0.6.5,<0.9.0",
-  "transformers==4.50.1",
-  "librosa",
-]
-
-[project.urls]
-Modelpage = "https://huggingface.co/MERaLiON/MERaLiON-2-10B" 
-Homepage = "https://huggingface.co/MERaLiON/MERaLiON-2-10B/tree/main/vllm_plugin_meralion2"
-Documentation = "https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/vllm_plugin_meralion2/readme.md"
-
-[project.entry-points."vllm.general_plugins"]
-register_dummy_model = "vllm_plugin_meralion2:register"
-
-[tool.setuptools.packages.find]
-where = ["src"]
-
-[tool.setuptools]
-package-dir = {"" = "src"}
-include-package-data = true

vllm_plugin_meralion2/readme.md

@@ -1,38 +0,0 @@
-## MERaLiON2 vLLM Plugin
-
-### Licence
-
-[MERaLiON-Public-Licence-v2](https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/MERaLiON-Public-Licence-v2.pdf)
-
-### Set up Environment
-
-This vLLM plugin for MERaLiON2 requires transformers version `4.50.1`. It supports vLLM version `0.6.5` ~ `0.7.3` (V0 engine), and `0.8.5` ~ `0.8.5.post1` (V1 engine). 
-
-```bash
-pip install transformers==4.50.1
-pip install vllm==0.6.5
-```
-
-Install the MERaLiON2 vLLM plugin.
-
-```bash
-pip install vllm-plugin-meralion2
-```
-
-It's strongly recommended to install flash-attn for better memory and gpu utilization. 
-
-```bash
-pip install flash-attn --no-build-isolation
-```
-
-### Offline Inference
-
-Refer to [offline_example.py](https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/vllm_plugin_meralion2/offline_example.py) for offline inference example.
-
-### OpenAI-compatible Serving
-
-Refer to [openai_serve_example.sh](https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/vllm_plugin_meralion2/openai_serve_example.sh) for openAI-compatible serving example.
-
-To call the server, you can refer to [openai_client_example.py](https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/vllm_plugin_meralion2/openai_client_example.py).
-
-Alternatively, you can try calling the server with curl, refer to [openai_client_curl.sh](https://huggingface.co/MERaLiON/MERaLiON-2-10B/blob/main/vllm_plugin_meralion2/openai_client_curl.sh).

vllm_plugin_meralion2/src/vllm_plugin_meralion2/__init__.py

@@ -1,46 +0,0 @@
-from typing import Optional
-
-from vllm.entrypoints.chat_utils import BaseMultiModalItemTracker
-
-from .transformers_utils.no_repeat_logits_processor import NoRepeatNGramLogitsProcessor
-
-
-_original_placeholder_str = BaseMultiModalItemTracker._placeholder_str
-
-
-def custom_placeholder_str(self, modality,
-                        current_count: int) -> Optional[str]:
-    hf_config = self._model_config.hf_config
-    model_type = hf_config.model_type
-
-    if modality == "audio" and model_type == "meralion2":
-        return "<SpeechHere>"
-
-    return _original_placeholder_str(self, modality=modality, current_count=current_count)
-
-
-def register():
-    import vllm
-    from vllm import ModelRegistry
-
-    v064_compatible_versions = ['0.6.5', '0.6.6', '0.6.6.post1', '0.7.0', '0.7.1', '0.7.2', '0.7.3']
-    v085_compatible_versions = ['0.8.5', '0.8.5.post1']
-    sorted_compatible_versions = sorted(v064_compatible_versions + v085_compatible_versions)
-
-    if vllm.__version__ in v064_compatible_versions:
-        from .vllm064_post1 import MERaLiON2ForConditionalGeneration
-    elif vllm.__version__ in v085_compatible_versions:
-        from .vllm085 import MERaLiON2ForConditionalGeneration
-    else:
-        raise Exception((
-            f"MERaLiON2 doesn't support vLLM version {vllm.__version__}."
-            f" Supported vLLM versions: {', '.join((sorted_compatible_versions))}"
-            ))
-
-    if "MERaLiON2ForConditionalGeneration" not in ModelRegistry.get_supported_archs():
-        ModelRegistry.register_model(
-            "MERaLiON2ForConditionalGeneration",
-            MERaLiON2ForConditionalGeneration
-            )
-        
-    vllm.entrypoints.chat_utils.BaseMultiModalItemTracker._placeholder_str = custom_placeholder_str

vllm_plugin_meralion2/src/vllm_plugin_meralion2/transformers_utils/configuration_meralion2.py

@@ -1,76 +0,0 @@
-"""MERaLiON2 model configuration"""
-
-from transformers import Gemma2Config, WhisperConfig
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class MERaLiON2Config(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`MERaLiON2ForConditionalGeneration`]. It is used to instantiate an
-    MERaLiON2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
-    with the defaults will yield a similar configuration to that of the MERaLiON2.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        audio_config (`Union[AutoConfig, dict]`,  *optional*, defaults to `CLIPVisionConfig`):
-            The config object or dictionary of the audio backbone.
-        text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `LlamaConfig`):
-            The config object or dictionary of the text backbone.
-        audio_token_index (`int`, *optional*, defaults to 151646):
-            The image token index to encode the image prompt.
-    """
-
-    model_type = "meralion2"
-    is_composition = False
-
-    def __init__(
-        self,
-        speech_config=None,
-        text_config=None,
-        speech_mlp_scale_factor=15,
-        speech_token_index=255999,
-        **kwargs,
-    ):
-        
-        if isinstance(speech_config, dict):
-            speech_config = WhisperConfig(**speech_config)
-        elif speech_config is None:
-            speech_config = WhisperConfig(
-                d_model=1280,
-                encoder_attention_heads=20,
-                encoder_ffn_dim=5120,
-                encoder_layerdrop=0.0,
-                encoder_layers=32,
-                num_mel_bins=128,
-                max_source_positions=1500,
-                scale_embedding=False,
-                activation_function="gelu",
-            )
-
-        self.speech_config = speech_config
-
-        if isinstance(text_config, dict):
-            text_config = Gemma2Config(**text_config)
-        elif text_config is None:
-            text_config = Gemma2Config()
-
-        self.text_config = text_config
-
-        self.speech_mlp_scale_factor = speech_mlp_scale_factor
-        self.speech_token_index = speech_token_index
-        
-        self.sliding_window = self.text_config.sliding_window
-        self.hidden_size = self.text_config.hidden_size
-        self.num_attention_heads = self.text_config.num_attention_heads
-        self.num_hidden_layers = self.text_config.num_hidden_layers
-        self.num_key_value_heads = self.text_config.num_key_value_heads
-        self.head_dim = self.text_config.head_dim
-        self.intermediate_size = self.text_config.intermediate_size
-        
-        super().__init__(**kwargs)

vllm_plugin_meralion2/src/vllm_plugin_meralion2/transformers_utils/modules.py

@@ -1,72 +0,0 @@
-"""Inference-only MERaLiON AudioLLM model compatible with HuggingFace weights."""
-from typing import Any, Optional, Set, Tuple, TypedDict, Union, List
-
-import torch
-import torch.nn as nn
-from transformers.utils.import_utils import is_torch_sdpa_available, is_flash_attn_2_available
-
-# === Audio Inputs === #
-class MERaLiON2Inputs(TypedDict):
-    input_features: torch.Tensor
-    """Shape: 
-    `(num_audios, num_mel_bins, 3000)`
-    """
-
-    feature_attention_mask: torch.Tensor
-    """Shape: `(num_audios, 3000)`
-    """
-
-
-# === Audio Encoder === #
-class MERaLiON2SpeechAudioAdaper(nn.Module):
-    def __init__(self, audio_hidden_size: int, text_hidden_size: int):
-        super(MERaLiON2SpeechAudioAdaper, self).__init__()
-        speech_mlp_scale_factor = 15
-
-        self.speech_mlp_scale_factor = speech_mlp_scale_factor
-        self.mlp_adapter = nn.Sequential(
-            nn.Linear(
-                in_features=audio_hidden_size * speech_mlp_scale_factor,
-                out_features=audio_hidden_size * 5,
-            ),
-            nn.SiLU(),
-        )
-
-        self.gate_proj = nn.Linear(
-                in_features=audio_hidden_size * 5,
-                out_features=audio_hidden_size * 5,
-            )
-        
-        self.pool_proj = nn.Linear(
-                in_features=audio_hidden_size * 5,
-                out_features=audio_hidden_size * 5,
-        )
-        self.act_fn = nn.SiLU()
-        self.out_proj = nn.Linear(
-                audio_hidden_size * 5,
-                text_hidden_size,
-        )
-
-    def forward(self, speech_embeds, **kwargs):
-        B, T, C = speech_embeds.shape
-        speech_embeds = self.mlp_adapter(
-            speech_embeds.reshape(
-                B,
-                T // self.speech_mlp_scale_factor,
-                C * self.speech_mlp_scale_factor,
-            )
-        )
-        speech_embeds = self.act_fn(self.gate_proj(speech_embeds)) * self.pool_proj(speech_embeds)
-        speech_embeds = self.out_proj(speech_embeds)
-        return speech_embeds
-    
-
-def autoset_attn_implementation_for_whisper(config):
-    _implementation = "eager"
-    if is_torch_sdpa_available():
-        _implementation = "sdpa"
-    if is_flash_attn_2_available():
-        _implementation = "flash_attention_2"
-
-    config._attn_implementation = _implementation
-    return config

vllm_plugin_meralion2/src/vllm_plugin_meralion2/transformers_utils/no_repeat_logits_processor.py

@@ -1,98 +0,0 @@
-from typing import Iterable, List
-
-import torch
-
-
-def _get_ngrams(ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int):
-    """
-    Assume ngram_size=2 and prev_input_ids=tensor([[40, 2883, 2712, 4346]]). The output of generated ngrams look like
-    this {(40,): [2883], (2883,): [2712], (2712,): [4346]}.
-
-    Args:
-        ngram_size (`int`):
-            The number sequential tokens taken as a group which may only occur once before being banned.
-        prev_input_ids (`torch.Tensor`):
-           Generated token ids for the current hypothesis.
-        num_hypos (`int`):
-            The number of hypotheses for which n-grams need to be generated.
-
-    Returns:
-        generated_ngrams (`dict`):
-            Dictionary of generated ngrams.
-    """
-    # Initialize an empty list of dictionaries, one for each hypothesis (index) in the range of num_hypos
-    generated_ngrams = [{} for _ in range(num_hypos)]
-    for idx in range(num_hypos):
-        gen_tokens = prev_input_ids[idx].tolist()
-        generated_ngram = generated_ngrams[idx]
-        # Loop through each n-gram of size ngram_size in the list of tokens (gen_tokens)
-        for ngram in zip(*[gen_tokens[i:] for i in range(ngram_size)]):
-            prev_ngram_tuple = tuple(ngram[:-1])
-            generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]
-    return generated_ngrams
-
-
-def _get_generated_ngrams(banned_ngrams, prev_input_ids, ngram_size, cur_len):
-    """
-    Determines the banned tokens for the current hypothesis based on previously generated n-grams.
-
-    Args:
-        banned_ngrams (`dict`):
-            A dictionary containing previously generated n-grams for each hypothesis.
-        prev_input_ids (`torch.Tensor`):
-            Generated token ids for the current hypothesis.
-        ngram_size (`int`):
-            The number sequential tokens taken as a group which may only occur once before being banned.
-        cur_len (`int`):
-            The current length of the token sequences for which the n-grams are being checked.
-
-    Returns:
-        List of tokens that are banned.
-    """
-    # Before decoding the next token, prevent decoding of ngrams that have already appeared
-    start_idx = cur_len + 1 - ngram_size
-    ngram_idx = tuple(prev_input_ids[start_idx:cur_len].tolist())
-    return banned_ngrams.get(ngram_idx, [])
-
-
-def _calc_banned_ngram_tokens(
-    ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int, cur_len: int
-) -> List[Iterable[int]]:
-    """Copied from fairseq for no_repeat_ngram in beam_search"""
-    if cur_len + 1 < ngram_size:
-        # return no banned tokens if we haven't generated no_repeat_ngram_size tokens yet
-        return [[] for _ in range(num_hypos)]
-    generated_ngrams = _get_ngrams(ngram_size, prev_input_ids, num_hypos)
-    banned_tokens = [
-        _get_generated_ngrams(generated_ngrams[hypo_idx], prev_input_ids[hypo_idx], ngram_size, cur_len)
-        for hypo_idx in range(num_hypos)
-    ]
-    return banned_tokens
-
-
-class NoRepeatNGramLogitsProcessor:
-
-    def __init__(self, ngram_size: int=6):
-        if not isinstance(ngram_size, int) or ngram_size <= 0:
-            raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}")
-        self.ngram_size = ngram_size
-
-    def __call__(self, prompt_tokens_ids: tuple, past_tokens_ids: tuple, scores: torch.FloatTensor) -> torch.FloatTensor:
-        # score: [B, vocab_size]
-        # input_ids: [B, cur_len]
-        input_ids = prompt_tokens_ids + past_tokens_ids
-        if len(input_ids) < self.ngram_size:
-            return scores
-
-        if len(scores.shape) == 1:
-            scores = scores.reshape(1, -1)
-
-        num_batch_hypotheses = scores.shape[0]
-        input_ids = torch.LongTensor(input_ids).reshape(num_batch_hypotheses, -1)
-        cur_len = input_ids.shape[-1]
-        scores_processed = scores.clone()
-        banned_batch_tokens = _calc_banned_ngram_tokens(self.ngram_size, input_ids, num_batch_hypotheses, cur_len)
-        for i, banned_tokens in enumerate(banned_batch_tokens):
-            scores_processed[i, banned_tokens] = -float("inf")
-
-        return scores_processed

vllm_plugin_meralion2/src/vllm_plugin_meralion2/transformers_utils/processing_meralion2.py

@@ -1,193 +0,0 @@
-"""Processor class for MERaLiON2."""
-
-from typing import List, Optional, Union
-
-import numpy as np
-
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.processing_utils import ProcessorMixin
-from transformers.tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput
-
-
-# copied from transformers.models.qwen2_audio.processing_qwen2_audio.Qwen2AudioProcessor
-class MERaLiON2Processor(ProcessorMixin):
-    r"""
-    Constructs a MERaLiON2 processor which wraps a whisper feature extractor and a gemma tokenizer into a single processor.
-
-    [`MERaLiON2Processor`] offers all the functionalities of [`WhisperFeatureExtractor`] and [`GemmaTokenizer`]. See the
-    [`~MERaLiON2Processor.__call__`] and [`~MERaLiON2Processor.decode`] for more information.
-
-    Args:
-        feature_extractor ([`WhisperFeatureExtractor`], *optional*):
-            The feature extractor is a required input.
-        tokenizer ([`GemmaTokenizer`], *optional*):
-            The tokenizer is a required input.
-        chat_template (`Optional[str]`, *optional*):
-                The Jinja template to use for formatting the conversation. If not provided, the default chat template
-                is used.
-    """
-
-    attributes = ["feature_extractor", "tokenizer"]
-    feature_extractor_class = "WhisperFeatureExtractor"
-    tokenizer_class = "AutoTokenizer"
-    valid_kwargs = [
-        "fixed_speech_embeds_length", 
-        "speech_token_index", 
-        "time_duration_limit", 
-        "whisper_chunk_size",
-        "do_normalize"
-    ]
-
-    def __init__(
-        self, 
-        feature_extractor=None, 
-        tokenizer=None, 
-        fixed_speech_embeds_length=100,
-        speech_token_index=255999,
-        time_duration_limit=300,
-        whisper_chunk_size=30,
-        do_normalize=True
-    ):
-        self.fixed_speech_embeds_length = fixed_speech_embeds_length
-        self.speech_token_index = speech_token_index
-        self.time_duration_limit = time_duration_limit
-        self.whisper_chunk_size = whisper_chunk_size
-        self.number_chunk_limit = self.time_duration_limit // self.whisper_chunk_size
-        self.do_normalize = do_normalize
-
-        super().__init__(feature_extractor, tokenizer)
-
-        self.speech_token = self.tokenizer.added_tokens_decoder[self.speech_token_index].content
-        self.feature_chunk_size = self.whisper_chunk_size * self.feature_extractor.sampling_rate
-
-    def _process_text(self, text: List[str], audio_number_chunks: np.ndarray):
-        pieces = []
-        for i, item in enumerate(text):
-            target_string = self.speech_token * self.fixed_speech_embeds_length * audio_number_chunks[i]
-            pieces.append(item.replace(self.speech_token, target_string))
-        return pieces
-
-    def _get_number_chunks(self, audios: List[np.ndarray]):
-        audio_lengths = np.array([_.shape[0] for _ in audios])
-        number_chunks = ((audio_lengths - 1) // self.feature_chunk_size) + 1
-        return np.clip(number_chunks, a_min=None, a_max=self.number_chunk_limit)
-
-    def _get_chunked_audios(self, audios: Union[np.ndarray, List[np.ndarray]]):
-        if isinstance(audios, np.ndarray):
-            audios = [audios]
-
-        audio_number_chunks = self._get_number_chunks(audios)
-        chunked_audios = []
-
-        for audio_idx, audio in enumerate(audios):
-            for cid in range(audio_number_chunks[audio_idx]):
-                chunked_audios.append(
-                    audio[cid * self.feature_chunk_size: (cid + 1) * self.feature_chunk_size]
-                )
-        return audio_number_chunks, chunked_audios
-
-    def __call__(
-        self,
-        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
-        audios: Union[np.ndarray, List[np.ndarray]] = None,
-        padding: Union[bool, str, PaddingStrategy] = True,
-        sampling_rate: Optional[int] = None,
-        do_normalize: Optional[bool] = None,
-        **kwargs,
-    ) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text`
-        and `kwargs` arguments to GemmaTokenizer's [`~GemmaTokenizer.__call__`] if `text` is not `None` to encode
-        the text. To prepare the audio(s), this method forwards the `audios` and `kwrags` arguments to
-        WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audios` is not `None`. Please refer to the doctsring
-        of the above two methods for more information.
-
-        Args:
-            text (`str`, `List[str]`):
-                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
-                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
-                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
-            audios (`np.ndarray`, `List[np.ndarray]`):
-                The audio or batch of audios to be prepared. Each audio can be a NumPy array.
-            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
-                Select a strategy to pad the returned sequences (according to the model's padding side and padding
-                index) among:
-                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
-                  sequence if provided).
-                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
-                  acceptable input length for the model if that argument is not provided.
-                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
-                  lengths).
-            sampling_rate (`int`, defaults to 16000):
-                The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
-            do_normalize (`bool`, defaults to `True`):
-                Whether or not to zero-mean unit-variance normalize the input. 
-                Normalizing can help to significantly improve the performance of the model.
-        """
-
-        if text is None:
-            raise ValueError("You need to specify either a `text` input to process.")
-        if not isinstance(text, list):
-            text = [text]
-        if not isinstance(audios, list):
-            audios = [audios]
-        if sampling_rate is None:
-            sampling_rate = self.feature_extractor.sampling_rate
-        if do_normalize is None:
-            do_normalize = self.do_normalize
-
-        for i, audio in enumerate(audios):
-            if audio.ndim > 1:
-                raise Exception(f"MERaLiON2 only accepts mono channel audio, {i+1}th audio have {audios[0].ndim} channels")
-        
-        inputs_dict = {}
-        
-        if audios is not None:
-            audio_number_chunks, chunked_audios = self._get_chunked_audios(audios)
-            text = self._process_text(text, audio_number_chunks)
-            
-            audio_inputs = self.feature_extractor(
-                chunked_audios, 
-                sampling_rate=sampling_rate, 
-                return_tensors="pt",
-                return_attention_mask=True, 
-                padding="max_length", 
-                do_normalize=self.do_normalize,
-            )
-            audio_inputs["feature_attention_mask"] = audio_inputs.pop(
-                "attention_mask"
-            )  # rename attention_mask to prevent conflicts later on
-            inputs_dict.update(audio_inputs)
-
-        text_input = self.tokenizer(
-            text=text,
-            return_tensors="pt",
-            add_special_tokens=False,
-            return_attention_mask=True,
-            padding=padding,
-        )   
-
-        inputs_dict["input_ids"] = text_input.input_ids
-        inputs_dict["attention_mask"] = text_input.attention_mask
-
-        return BatchFeature(data={**inputs_dict})
-
-    def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to GemmaTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
-        refer to the docstring of this method for more information.
-        """
-        return self.tokenizer.batch_decode(*args, **kwargs)
-
-    def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to GemmaTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
-        the docstring of this method for more information.
-        """
-        return self.tokenizer.decode(*args, **kwargs)
-
-    @property
-    def model_input_names(self):
-        tokenizer_input_names = self.tokenizer.model_input_names
-        feature_extractor_input_names = self.feature_extractor.model_input_names
-        return list(dict.fromkeys(tokenizer_input_names + feature_extractor_input_names + ["feature_attention_mask"]))

vllm_plugin_meralion2/src/vllm_plugin_meralion2/vllm064_post1.py

@@ -1,454 +0,0 @@
-"""Inference-only MERaLiON AudioLLM model compatible with HuggingFace weights."""
-from functools import lru_cache
-from typing import Iterable, List, Mapping, Optional, Tuple, Union
-
-import librosa
-import numpy as np
-import torch
-import torch.nn as nn
-
-from vllm.attention import AttentionMetadata
-from vllm.config import VllmConfig
-from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,
-                         InputContext, token_inputs)
-from vllm.logger import init_logger
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader, maybe_remap_kv_scale_name)
-from vllm.model_executor.models.gemma2 import Gemma2Model
-from vllm.model_executor.sampling_metadata import SamplingMetadata
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs
-from vllm.multimodal.utils import consecutive_placeholder_ranges
-from vllm.sequence import IntermediateTensors, SequenceData
-from transformers.models.whisper.modeling_whisper import WhisperEncoder
-
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsLoRA, SupportsPP
-from vllm.model_executor.models.utils import maybe_prefix
-
-from .transformers_utils.modules import (autoset_attn_implementation_for_whisper,
-    MERaLiON2Inputs, MERaLiON2SpeechAudioAdaper)
-
-
-logger = init_logger(__name__)
-
-
-# gemma2 ties word embedding by default
-_KEYS_TO_MODIFY_MAPPING = {
-    "text_decoder.model": "model",
-}
-
-# === Constants === #
-DEFAULT_SAMPLE_RATE = 16000
-FEATURE_CHUNK_SIZE = DEFAULT_SAMPLE_RATE * 30
-OUTPUT_CHUNK_SIZE = 100
-MAX_NUMBER_CHUNKS = 8
-
-
-def dummy_data_for_meralion(ctx: InputContext, seq_len: int,
-                               mm_counts: Mapping[str, int]):
-    num_audios = mm_counts["audio"]
-    max_tokens_per_audio = get_max_meralion_audio_tokens(ctx)
-    max_llm_audio_tokens = max_tokens_per_audio * num_audios
-    if seq_len - max_llm_audio_tokens - 2 < 0:
-        raise RuntimeError(
-            f"MERaLiON-AudioLLM cannot process {num_audios} audios in a prompt, "
-            "please increase max_model_len or reduce audio limit by "
-            "--limit-mm-per-prompt.")
-
-    speech_token_index = ctx.model_config.hf_config.speech_token_index
-
-    dummy_seqdata = SequenceData.from_prompt_token_counts(
-        (speech_token_index, max_llm_audio_tokens),
-        (0, seq_len - max_llm_audio_tokens),
-    )
-    dummy_audio = np.full((max_llm_audio_tokens * 15 * 2 * 160, ), 0.)
-    return DummyData(
-        dummy_seqdata, {"audio": [(dummy_audio, DEFAULT_SAMPLE_RATE)] * num_audios}, {
-            "audio":
-            consecutive_placeholder_ranges(num_items=num_audios,
-                                           item_size=max_tokens_per_audio)
-        })
-
-
-def get_processor(
-    processor_name: str,
-    *args,
-    trust_remote_code: bool = True,
-    **kwargs,
-):
-    """Gets a processor for the given model name via HuggingFace.
-
-    Derived from `vllm.transformers_utils.image_processor.get_image_processor`.
-    """
-    # don't put this import at the top level
-    # it will call torch.cuda.device_count()
-    from transformers import AutoProcessor
-
-    try:
-        processor = AutoProcessor.from_pretrained(
-            processor_name,
-            *args,
-            trust_remote_code=trust_remote_code,
-            **kwargs)
-    except ValueError as e:
-        # If the error pertains to the processor class not existing or not
-        # currently being imported, suggest using the --trust-remote-code flag.
-        # Unlike AutoTokenizer, AutoProcessor does not separate such errors
-        if not trust_remote_code:
-            err_msg = (
-                "Failed to load the processor. If the processor is "
-                "a custom processor not yet available in the HuggingFace "
-                "transformers library, consider setting "
-                "`trust_remote_code=True` in LLM or using the "
-                "`--trust-remote-code` flag in the CLI.")
-            raise RuntimeError(err_msg) from e
-        else:
-            raise e
-
-    return processor
-
-
-cached_get_processor = lru_cache(get_processor)
-
-
-def _get_number_chunks(audios: List[np.ndarray]):
-    audio_lengths = np.array([_.shape[0] for _ in audios])
-    number_chunks = ((audio_lengths - 1) // FEATURE_CHUNK_SIZE) + 1
-    return np.clip(number_chunks, a_min=None, a_max=MAX_NUMBER_CHUNKS)
-
-
-def _get_feat_extract_output_lengths(audios: List[np.ndarray]):
-    return _get_number_chunks(audios) * OUTPUT_CHUNK_SIZE
-
-
-def _get_chunked_audios(audios: List[np.ndarray]):
-    audio_number_chunks = _get_number_chunks(audios)
-    chunked_resampled_audios = [] 
-
-    for audio_idx, audio in enumerate(audios):
-        for cid in range(audio_number_chunks[audio_idx]):
-            chunked_resampled_audios.append(
-                audio[cid * FEATURE_CHUNK_SIZE: (cid + 1) * FEATURE_CHUNK_SIZE]
-            )
-    return chunked_resampled_audios
-
-
-def _maybe_resample_audio(audio, orig_sample_rate, target_sample_rate):
-    if orig_sample_rate != target_sample_rate:
-        return librosa.resample(
-            audio,
-            orig_sr=orig_sample_rate,
-            target_sr=target_sample_rate
-            )
-    return audio
-
-
-def get_max_meralion_audio_tokens(ctx: InputContext) -> int:
-    """
-    The max number of tokens after speech audio adapter.
-    """
-    return MAX_NUMBER_CHUNKS * OUTPUT_CHUNK_SIZE
-
-
-def input_processor_for_meralion(
-        ctx: InputContext, inputs: DecoderOnlyInputs) -> DecoderOnlyInputs:
-    multi_modal_data = inputs.get("multi_modal_data")
-    if multi_modal_data is None or "audio" not in multi_modal_data:
-        return inputs
-
-    audios = multi_modal_data["audio"]
-    if not isinstance(audios, list):
-        audios = [audios]
-
-    if len(audios) == 0:
-        return inputs
-
-    processor = cached_get_processor(ctx.model_config.model)
-
-    resampled_audios = [
-        librosa.resample(audio,
-                         orig_sr=sampling_rate,
-                         target_sr=processor.feature_extractor.sampling_rate)
-        for audio, sampling_rate in audios
-    ]
-
-    audio_output_lengths = _get_feat_extract_output_lengths(resampled_audios)
-    speech_token_index = ctx.model_config.hf_config.speech_token_index
-
-    input_ids = inputs['prompt_token_ids']
-
-    new_input_ids = []
-    audio_num = input_ids.count(speech_token_index)
-    assert len(audio_output_lengths) == audio_num, \
-        (f'The text input contains {audio_num} audio tokens, '
-         f'but {len(audio_output_lengths)} audios provided')
-    start = 0
-    for audio_idx in range(audio_num):
-        end = input_ids.index(speech_token_index, start)
-        new_input_ids.extend(input_ids[start:end])  # text part
-
-        new_input_ids.extend([speech_token_index] * 
-                             audio_output_lengths[audio_idx])
-        start = end + 1
-    new_input_ids.extend(input_ids[start:])
-
-    return token_inputs(
-        prompt_token_ids=new_input_ids,
-        prompt=inputs.get('prompt'),
-        multi_modal_data=multi_modal_data,
-    )
-
-
-def input_mapper_for_meralion(
-    ctx: InputContext,
-    multi_modal_data: Union[np.ndarray, List[np.ndarray]],
-) -> MultiModalKwargs:
-    """Input mapper for MERaLiON-AudioLLM."""
-    if not isinstance(multi_modal_data, list):
-        multi_modal_data = [multi_modal_data]
-
-    if len(multi_modal_data) == 0:
-        return MultiModalKwargs()
-
-    processor = cached_get_processor(ctx.model_config.model)
-    audio_feature_extractor = processor.feature_extractor
-    if audio_feature_extractor is None:
-        raise RuntimeError(
-            "No HuggingFace audio_feature_extractor is available "
-            "to process the audio object")
-
-    try:
-        target_sample_rate = processor.feature_extractor.sampling_rate
-
-        resampled_audios = [
-            _maybe_resample_audio(
-                audio=audio,
-                orig_sample_rate=sampling_rate, 
-                target_sample_rate=target_sample_rate, 
-                )
-            for audio, sampling_rate in multi_modal_data
-        ]
-
-        resampled_audios = _get_chunked_audios(resampled_audios)
-
-        batch_data = audio_feature_extractor(resampled_audios,
-                                             sampling_rate=target_sample_rate,
-                                             return_attention_mask=True,
-                                             padding="max_length",
-                                             return_tensors="pt",
-                                             do_normalize=True).data
-        batch_data["feature_attention_mask"] = batch_data.pop("attention_mask")
-    except Exception:
-        logger.error("Failed to process audio (%s)", multi_modal_data)
-        raise
-
-    return MultiModalKwargs(batch_data)
-
-
-@INPUT_REGISTRY.register_dummy_data(dummy_data_for_meralion)
-@INPUT_REGISTRY.register_input_processor(input_processor_for_meralion)
-@MULTIMODAL_REGISTRY.register_input_mapper("audio",
-                                           input_mapper_for_meralion)
-@MULTIMODAL_REGISTRY.register_max_multimodal_tokens(
-    "audio", get_max_meralion_audio_tokens)
-class MERaLiON2ForConditionalGeneration(nn.Module, SupportsMultiModal,
-                                       SupportsLoRA, SupportsPP):
-    packed_modules_mapping = {
-        "qkv_proj": [
-            "q_proj",
-            "k_proj",
-            "v_proj",
-        ],
-        "gate_up_proj": [
-            "gate_proj",
-            "up_proj",
-        ],
-    }
-
-    # LoRA specific attributes
-    supported_lora_modules = [
-        "qkv_proj",
-        "o_proj",
-        "gate_up_proj",
-        "down_proj",
-    ]
-
-    embedding_modules = {}
-    embedding_padding_modules = []
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        multimodal_config = vllm_config.model_config.multimodal_config
-
-        self.config = config
-        self.multimodal_config = multimodal_config
-
-        config.speech_config = \
-            autoset_attn_implementation_for_whisper(config.speech_config)
-        self.speech_encoder = WhisperEncoder(config.speech_config)
-        self.ln_speech = nn.LayerNorm(config.speech_config.d_model)
-        self.speech_audio_adapter = MERaLiON2SpeechAudioAdaper(
-            config.speech_config.d_model, config.text_config.hidden_size)
-
-        self.quant_config = quant_config
-
-        self.model = Gemma2Model(
-            vllm_config=vllm_config.with_hf_config(config.text_config),
-            prefix=maybe_prefix(prefix, "model"))
-        self.unpadded_vocab_size = config.text_config.vocab_size
-        if config.text_config.tie_word_embeddings:
-            self.lm_head = self.model.embed_tokens
-        else:
-            self.lm_head = ParallelLMHead(config.text_config.vocab_size,
-                                          config.text_config.hidden_size,
-                                          quant_config=quant_config)
-        logit_scale = getattr(config, "logit_scale", 1.0)
-        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
-                                                config.text_config.vocab_size,
-                                                logit_scale)
-
-        self.sampler = get_sampler()
-        self.make_empty_intermediate_tensors = (
-            self.model.make_empty_intermediate_tensors)
-
-    def _validate_and_reshape_mm_tensor(self,
-                                        mm_input: Union[torch.Tensor,
-                                                        List[torch.Tensor]],
-                                        name: str) -> torch.Tensor:
-        if not isinstance(mm_input, (torch.Tensor, list)):
-            raise ValueError(f"Incorrect type of {name}. "
-                             f"Got type: {type(mm_input)}")
-        if isinstance(mm_input, torch.Tensor):
-            return torch.concat(list(mm_input))
-        else:
-            return torch.concat(mm_input)
-
-    def _parse_and_validate_audio_input(
-            self, **kwargs: object) -> Optional[MERaLiON2Inputs]:
-        input_features = kwargs.pop('input_features', None)
-        feature_attention_mask = kwargs.pop('feature_attention_mask', None)
-        if input_features is None:
-            return None
-        input_features = self._validate_and_reshape_mm_tensor(
-            input_features, 'input_features')
-        feature_attention_mask = self._validate_and_reshape_mm_tensor(
-            feature_attention_mask, 'feature_attention_mask')
-        if not isinstance(input_features, (torch.Tensor, list)):
-            raise ValueError("Incorrect type of audio input features. "
-                             f"Got type: {type(input_features)}")
-        return MERaLiON2Inputs(input_features=input_features,
-                                feature_attention_mask=feature_attention_mask)
-
-    def _process_audio_input(self,
-                             audio_input: MERaLiON2Inputs) -> torch.Tensor:
-
-        input_features = audio_input["input_features"].to(self.speech_encoder.dtype)
-        feature_attention_mask = audio_input["feature_attention_mask"]
-
-        audio_outputs = self.speech_encoder(input_features,
-                                         attention_mask=feature_attention_mask)
-        audio_features = audio_outputs.last_hidden_state
-        audio_features = self.ln_speech(audio_features)
-        audio_features = self.speech_audio_adapter(audio_features)
-        audio_features = audio_features.view(-1, audio_features.size(-1))
-        
-        return audio_features
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        kv_caches: List[torch.Tensor],
-        attn_metadata: AttentionMetadata,
-        intermediate_tensors: Optional[IntermediateTensors] = None,
-        **kwargs: object,
-    ) -> Union[torch.Tensor, IntermediateTensors]:
-        if intermediate_tensors is not None:
-            input_ids = None
-            inputs_embeds = None
-        else:
-            audio_input = self._parse_and_validate_audio_input(**kwargs)
-
-            if audio_input is None:
-                inputs_embeds = None
-            else:
-                inputs_embeds = self.model.embed_tokens(input_ids)
-                processed_audio_features = self._process_audio_input(audio_input)
-                # merge llm embeddings and audio features
-                mask = (input_ids == self.config.speech_token_index)
-                inputs_embeds[mask, :] = processed_audio_features
-
-                input_ids = None
-
-        hidden_states = self.model(
-            input_ids=input_ids,
-            positions=positions,
-            kv_caches=kv_caches,
-            attn_metadata=attn_metadata,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
-        )
-        return hidden_states
-
-    def compute_logits(self, hidden_states: torch.Tensor,
-                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
-        logits = self.logits_processor(self.lm_head, hidden_states,
-                                       sampling_metadata)
-        return logits
-
-    def sample(
-        self,
-        logits: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
-    ) -> Optional[SamplerOutput]:
-        next_tokens = self.sampler(logits, sampling_metadata)
-        return next_tokens
-
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-            if (self.config.text_config.tie_word_embeddings
-                    and "lm_head.weight" in name):
-                continue
-            for key_to_modify, new_key in _KEYS_TO_MODIFY_MAPPING.items():
-                if key_to_modify in name:
-                    name = name.replace(key_to_modify, new_key)
-            for (param_name, weight_name, shard_id) in stacked_params_mapping:
-                if weight_name not in name or 'speech_' in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                # Remapping the name of FP8 kv-scale.
-                name = maybe_remap_kv_scale_name(name, params_dict)
-                if name is None:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
-                weight_loader(param, loaded_weight)

vllm_plugin_meralion2/src/vllm_plugin_meralion2/vllm085.py

@@ -1,333 +0,0 @@
-"""Inference-only MERaLiON AudioLLM model compatible with HuggingFace weights."""
-from collections.abc import Iterable, Mapping, Sequence
-from typing import Any, Optional, Set, Tuple, TypedDict, Union, List
-
-import math
-import torch
-import torch.nn as nn
-from transformers import BatchFeature
-from transformers.models.whisper.modeling_whisper import WhisperEncoder
-from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor
-
-from vllm.config import VllmConfig
-from vllm.model_executor.sampling_metadata import SamplingMetadata
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs
-from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
-                                    MultiModalKwargs)
-from vllm.multimodal.parse import (AudioProcessorItems, MultiModalDataItems,
-                                   MultiModalDataParser)
-from vllm.multimodal.processing import (BaseMultiModalProcessor,
-                                        BaseProcessingInfo, PromptReplacement,
-                                        PromptUpdate, PromptUpdateDetails)
-from vllm.multimodal.profiling import BaseDummyInputsBuilder
-from vllm.sequence import IntermediateTensors
-from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
-from vllm.model_executor.models.utils import (AutoWeightsLoader, init_vllm_registered_model,
-                    maybe_prefix, merge_multimodal_embeddings)
-
-from .transformers_utils.processing_meralion2 import MERaLiON2Processor
-from .transformers_utils.configuration_meralion2 import MERaLiON2Config
-from .transformers_utils.modules import (autoset_attn_implementation_for_whisper,
-    MERaLiON2Inputs, MERaLiON2SpeechAudioAdaper)
-
-
-class MERaLiON2ProcessingInfo(BaseProcessingInfo):
-
-    def get_hf_config(self):
-        return self.ctx.get_hf_config(MERaLiON2Config)
-
-    def get_hf_processor(
-        self,
-        *,
-        # Ignored in initialization
-        sampling_rate: Optional[int] = None,
-        **kwargs: object,
-    ) -> MERaLiON2Processor:
-        return self.ctx.get_hf_processor(MERaLiON2Processor, **kwargs)
-
-    def get_feature_extractor(
-        self,
-        *,
-        # Ignored in initialization
-        sampling_rate: Optional[int] = None,
-    ) -> WhisperFeatureExtractor:
-        hf_processor = self.get_hf_processor(sampling_rate=sampling_rate)
-        feature_extractor = hf_processor.feature_extractor  # type: ignore
-        assert isinstance(feature_extractor, WhisperFeatureExtractor)
-        return feature_extractor
-
-    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
-        return {"audio": None}
-    
-
-class MERaLiON2DummyInputsBuilder(
-        BaseDummyInputsBuilder[MERaLiON2ProcessingInfo]):
-
-    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
-        num_audios = mm_counts.get("audio", 0)
-
-        hf_processor = self.info.get_hf_processor()
-        audio_token = hf_processor.speech_token
-
-        return audio_token * num_audios
-
-    def get_dummy_mm_data(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-    ) -> MultiModalDataDict:
-        processor = self.info.get_hf_processor()
-        feature_extractor = self.info.get_feature_extractor()
-
-        # This is to specify the audio length
-        num_chunk_limit = processor.number_chunk_limit
-        sampling_rate = feature_extractor.sampling_rate
-        audio_len = num_chunk_limit * feature_extractor.chunk_length * sampling_rate 
-        num_audios = mm_counts.get("audio", 0)
-
-        return {
-            "audio":
-            self._get_dummy_audios(length=audio_len, num_audios=num_audios)
-        }
-    
-
-class MERaLiON2MultiModalProcessor(
-        BaseMultiModalProcessor[MERaLiON2ProcessingInfo]):
-
-    def _get_data_parser(self) -> MultiModalDataParser:
-        feature_extractor = self.info.get_feature_extractor()
-        return MultiModalDataParser(target_sr=feature_extractor.sampling_rate)
-
-    def _call_hf_processor(
-        self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, Any],
-    ) -> BatchFeature:
-        # Text-only input not supported in composite processor
-        if not mm_data.get("audios", []):
-            prompt_ids = self.info.get_tokenizer().encode(prompt)
-            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
-            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
-
-        processor = self.info.get_hf_processor()
-        feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
-        speech_token_id = getattr(processor, "speech_token_index", 255999)
-        output_chunk_size = getattr(processor, "fixed_speech_embeds_length", 100)
-
-        mm_kwargs = dict(
-            **mm_kwargs,
-            sampling_rate=feature_extractor.sampling_rate,
-        )
-
-        results = super()._call_hf_processor(
-            prompt=prompt,
-            mm_data=mm_data,
-            mm_kwargs=mm_kwargs,
-        )
-
-        chunk_sizes = (results["input_ids"] == speech_token_id).sum(axis=1) // output_chunk_size
-        splitted_input_features = torch.split(results["input_features"], chunk_sizes.tolist())
-        splitted_feature_attention_mask = torch.split(results["feature_attention_mask"], chunk_sizes.tolist())
-
-        results["input_features"] = splitted_input_features
-        results["feature_attention_mask"] = splitted_feature_attention_mask
-        return results
-        
-    def _get_mm_fields_config(
-        self,
-        hf_inputs: BatchFeature,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> Mapping[str, MultiModalFieldConfig]:
-        return dict(
-            input_features=MultiModalFieldConfig.batched("audio"),
-            feature_attention_mask=MultiModalFieldConfig.batched("audio"),
-        )
-
-    def _get_prompt_updates(
-        self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        out_mm_kwargs: MultiModalKwargs,
-    ) -> Sequence[PromptUpdate]:
-        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-
-        speech_token = getattr(processor, "audio_token", "<SpeechHere>")
-        speech_token_id = getattr(processor, "speech_token_index", 255999)
-        output_chunk_size = getattr(processor, "fixed_speech_embeds_length", 100)
-        feature_chunk_size = getattr(processor, "feature_chunk_size", 30 * 16000)
-
-        def get_replacement_meralion2_audio(item_idx: int):
-            audios = mm_items.get_items("audio", AudioProcessorItems)
-            audio_length = audios.get_audio_length(item_idx)
-            number_chunks = ((audio_length - 1) // feature_chunk_size) + 1
-            speech_tokens = [speech_token_id] * number_chunks * output_chunk_size
-
-            return PromptUpdateDetails.select_token_id(
-                speech_tokens,
-                embed_token_id=speech_token_id,
-            )
-
-        return [
-            PromptReplacement(
-                modality="audio",
-                target=speech_token,
-                replacement=get_replacement_meralion2_audio,
-            )
-        ]
-    
-
-@MULTIMODAL_REGISTRY.register_processor(
-    MERaLiON2MultiModalProcessor,
-    info=MERaLiON2ProcessingInfo,
-    dummy_inputs=MERaLiON2DummyInputsBuilder)
-class MERaLiON2ForConditionalGeneration(nn.Module, SupportsMultiModal,
-                                         SupportsPP):
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        multimodal_config = vllm_config.model_config.multimodal_config
-        self.config = config
-        self.multimodal_config = multimodal_config
-
-        config.speech_config = \
-            autoset_attn_implementation_for_whisper(config.speech_config)
-        self.speech_encoder = WhisperEncoder(config.speech_config)
-        self.ln_speech = nn.LayerNorm(config.speech_config.d_model)
-        self.speech_audio_adapter = MERaLiON2SpeechAudioAdaper(
-            config.speech_config.d_model, config.text_config.hidden_size)
-
-        self.quant_config = quant_config
-
-        self.text_decoder = init_vllm_registered_model(
-            vllm_config=vllm_config,
-            hf_config=config.text_config,
-            prefix=maybe_prefix(prefix, "text_decoder"),
-            architectures=["Gemma2ForCausalLM"],
-        )
-
-        self.make_empty_intermediate_tensors = (
-            self.text_decoder.make_empty_intermediate_tensors)
-
-    def _validate_and_reshape_mm_tensor(self,
-                                        mm_input: Union[torch.Tensor,
-                                                        List[torch.Tensor]],
-                                        name: str) -> torch.Tensor:
-        if not isinstance(mm_input, (torch.Tensor, list)):
-            raise ValueError(f"Incorrect type of {name}. "
-                             f"Got type: {type(mm_input)}")
-        
-        input_ids_lst = mm_input.size()
-        with open("log.txt", "a") as f:
-            f.write(",".join([str(e) for e in input_ids_lst]) + "\n")
-
-        if isinstance(mm_input, torch.Tensor):
-            result = torch.concat(list(mm_input))            
-        else:
-            result = torch.concat(mm_input)
-
-        flattened_result = result.view(-1, result.size(-2), result.size(-1))
-        return flattened_result
-
-    def _parse_and_validate_audio_input(
-            self, **kwargs: object) -> Optional[MERaLiON2Inputs]:
-        input_features = kwargs.pop('input_features', None)
-        feature_attention_mask = kwargs.pop('feature_attention_mask', None)
-
-        if input_features is None:
-            return None
-        input_features = self._validate_and_reshape_mm_tensor(
-            input_features, 'input_features')
-        feature_attention_mask = self._validate_and_reshape_mm_tensor(
-            feature_attention_mask, 'feature_attention_mask')
-        if not isinstance(input_features, (torch.Tensor, list)):
-            raise ValueError("Incorrect type of audio input features. "
-                             f"Got type: {type(input_features)}")
-        return MERaLiON2Inputs(input_features=input_features,
-                                feature_attention_mask=feature_attention_mask)
-
-    def _process_audio_input(self,
-                             audio_input: MERaLiON2Inputs) -> torch.Tensor:
-        
-        input_features = audio_input["input_features"].to(self.speech_encoder.dtype)
-        feature_attention_mask = audio_input["feature_attention_mask"]
-
-        audio_outputs = self.speech_encoder(input_features,
-                                         attention_mask=feature_attention_mask)
-        audio_features = audio_outputs.last_hidden_state
-        audio_features = self.ln_speech(audio_features)
-        audio_features = self.speech_audio_adapter(audio_features)
-        audio_features = audio_features.view(-1, audio_features.size(-1))
-        
-        return audio_features
-    
-    def get_language_model(self) -> torch.nn.Module:
-        return self.text_decoder
-    
-    def get_multimodal_embeddings(
-            self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
-        
-        if "input_features" not in kwargs:
-            return None
-        
-        input_features = kwargs["input_features"]
-        audio_input = self._parse_and_validate_audio_input(**kwargs)        
-        if isinstance(input_features, torch.Tensor):
-            input_features = list(input_features)
-        
-        audio_lengths = [math.prod(audio.shape[:-2]) * 100 for audio in input_features]
-        masked_audio_features = self._process_audio_input(audio_input)
-        masked_audio_features = torch.split(masked_audio_features, audio_lengths)
-        return masked_audio_features
-    
-    def get_input_embeddings(
-        self,
-        input_ids: torch.Tensor,
-        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
-    ) -> torch.Tensor:
-        inputs_embeds = self.text_decoder.get_input_embeddings(input_ids)
-        if multimodal_embeddings is not None:
-            inputs_embeds = merge_multimodal_embeddings(
-                input_ids, inputs_embeds, multimodal_embeddings,
-                self.config.speech_token_index)
-        return inputs_embeds
-
-    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:
-            multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
-            inputs_embeds = self.get_input_embeddings(input_ids,
-                                                      multimodal_embeddings)
-            input_ids = None
-
-        hidden_states = self.text_decoder.model(input_ids,
-                                                  positions,
-                                                  intermediate_tensors,
-                                                  inputs_embeds=inputs_embeds)
-        return hidden_states
-
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
-    ) -> Optional[torch.Tensor]:
-        return self.text_decoder.compute_logits(hidden_states,
-                                                  sampling_metadata)
-    
-    def load_weights(self, weights: Iterable[Tuple[str,
-                                                   torch.Tensor]]) -> Set[str]:
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)