| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright 2025 Horizon team, Xiaomi MiLM Plus. |
| 4 | # Copyright 2024 The Qwen team. |
| 5 | # Copyright 2023 The vLLM team. |
| 6 | # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. |
| 7 | # |
| 8 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 9 | # and OPT implementations in this library. It has been modified from its |
| 10 | # original forms to accommodate minor architectural differences compared |
| 11 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 12 | # |
| 13 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 14 | # you may not use this file except in compliance with the License. |
| 15 | # You may obtain a copy of the License at |
| 16 | # |
| 17 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 18 | # |
| 19 | # Unless required by applicable law or agreed to in writing, software |
| 20 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 21 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 22 | # See the License for the specific language governing permissions and |
| 23 | # limitations under the License. |
| 24 | """Inference-only MiDashengLM model compatible with HuggingFace weights.""" |
| 25 | |
| 26 | import collections |
| 27 | import collections.abc |
| 28 | from collections.abc import Callable, Iterable, Mapping, Sequence |
| 29 | from typing import Annotated, TypeAlias, cast |
| 30 | |
| 31 | import numpy as np |
| 32 | import torch |
| 33 | import torch.nn as nn |
| 34 | import torchaudio.functional as F |
| 35 | from torch.nn.functional import scaled_dot_product_attention |
| 36 | from transformers import BatchFeature |
| 37 | |
| 38 | from vllm.config import VllmConfig |
| 39 | from vllm.config.multimodal import MultiModalDummyOptions |
| 40 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 41 | from vllm.inputs import MultiModalDataDict |
| 42 | from vllm.model_executor.layers.activation import get_act_fn |
| 43 | from vllm.model_executor.layers.conv import Conv2dLayer |
| 44 | from vllm.model_executor.layers.linear import ( |
| 45 | ColumnParallelLinear, |
| 46 | QKVParallelLinear, |
| 47 | RowParallelLinear, |
| 48 | ) |
| 49 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 50 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 51 | from vllm.multimodal.inputs import ( |
| 52 | MultiModalFieldConfig, |
| 53 | MultiModalKwargsItems, |
| 54 | ) |
| 55 | from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser |
| 56 | from vllm.multimodal.processing import ( |
| 57 | BaseDummyInputsBuilder, |
| 58 | BaseMultiModalProcessor, |
| 59 | BaseProcessingInfo, |
| 60 | PromptReplacement, |
| 61 | PromptUpdate, |
| 62 | PromptUpdateDetails, |
| 63 | ) |
| 64 | from vllm.multimodal.processing.processor import HFMultiModalInputs |
| 65 | from vllm.sequence import IntermediateTensors |
| 66 | from vllm.transformers_utils.configs.midashenglm import DashengConfig |
| 67 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 68 | |
| 69 | from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP |
| 70 | from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix |
| 71 | |
| 72 | _Tuple2: TypeAlias = int | tuple[int, int] | Sequence[int] |
| 73 | |
| 74 | |
| 75 | def _resolve_tuple2(x: _Tuple2) -> tuple[int, int]: |
| 76 | if isinstance(x, collections.abc.Sequence): |
| 77 | assert len(x) == 2, ( |
| 78 | f"Expected a sequence of length 2, got {x} with length {len(x)}" |
| 79 | ) |
| 80 | return cast(tuple[int, int], tuple(x)) |
| 81 | return (x, x) |
| 82 | |
| 83 | |
| 84 | def calculate_mel_frames_dasheng( |
| 85 | audio_length_samples: int, |
| 86 | n_fft: int = 512, |
| 87 | hop_size: int = 160, |
| 88 | dasheng_subsampling: int = 4, |
| 89 | center=True, |
| 90 | model_subsampling: int = 5, |
| 91 | ) -> int: |
| 92 | """Calculate the number of Mel-spectrogram frames.""" |
| 93 | if center: |
| 94 | audio_length_samples = audio_length_samples + n_fft |
| 95 | |
| 96 | return ( |
| 97 | int(1 + ((audio_length_samples - n_fft) / hop_size)) |
| 98 | // dasheng_subsampling |
| 99 | // model_subsampling |
| 100 | ) |
| 101 | |
| 102 | |
| 103 | class AudioPatchEmbed(nn.Module): |
| 104 | def __init__( |
| 105 | self, |
| 106 | input_size: _Tuple2 = 64, |
| 107 | patch_size: _Tuple2 = 16, |
| 108 | patch_stride: _Tuple2 = 16, |
| 109 | in_chans: int = 1, |
| 110 | embed_dim: int = 768, |
| 111 | norm_layer: Callable | None = None, |
| 112 | flatten: bool = False, |
| 113 | ): |
| 114 | super().__init__() |
| 115 | self.input_size = _resolve_tuple2(input_size) |
| 116 | self.patch_size = _resolve_tuple2(patch_size) |
| 117 | self.patch_stride = _resolve_tuple2(patch_stride) |
| 118 | self.grid_size = ( |
| 119 | self.input_size[0] // self.patch_stride[0], |
| 120 | self.input_size[1] // self.patch_stride[1], |
| 121 | ) |
| 122 | self.num_patches = self.grid_size[0] * self.grid_size[1] |
| 123 | self.flatten = flatten |
| 124 | |
| 125 | self.proj = Conv2dLayer( |
| 126 | in_chans, |
| 127 | embed_dim, |
| 128 | kernel_size=self.patch_size, |
| 129 | stride=self.patch_stride, |
| 130 | ) |
| 131 | self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() |
| 132 | |
| 133 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 134 | x = self.proj(x) |
| 135 | if self.flatten: |
| 136 | x = torch.permute( |
| 137 | torch.flatten(x, 2, 3), (0, 2, 1) |
| 138 | ) # rearrange(x, "b c f t -> b (f t) c") |
| 139 | x = self.norm(x) |
| 140 | return x |
| 141 | |
| 142 | |
| 143 | class LayerScale(nn.Module): |
| 144 | def __init__(self, dim, init_values=1e-5, inplace=False): |
| 145 | super().__init__() |
| 146 | self.inplace = inplace |
| 147 | self.gamma = nn.Parameter(init_values * torch.ones(dim)) |
| 148 | |
| 149 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 150 | return x.mul_(self.gamma) if self.inplace else x * self.gamma |
| 151 | |
| 152 | |
| 153 | class DashengMlp(nn.Module): |
| 154 | def __init__( |
| 155 | self, |
| 156 | in_features: int, |
| 157 | hidden_features: int | None = None, |
| 158 | out_features: int | None = None, |
| 159 | quant_config: QuantizationConfig | None = None, |
| 160 | prefix: str = "", |
| 161 | ): |
| 162 | super().__init__() |
| 163 | out_features = out_features or in_features |
| 164 | hidden_features = hidden_features or in_features |
| 165 | self.fc1 = ColumnParallelLinear( |
| 166 | input_size=in_features, |
| 167 | output_size=hidden_features, |
| 168 | quant_config=quant_config, |
| 169 | prefix=f"{prefix}.fc1", |
| 170 | ) |
| 171 | self.act = get_act_fn("gelu") |
| 172 | self.fc2 = RowParallelLinear( |
| 173 | input_size=hidden_features, |
| 174 | output_size=out_features, |
| 175 | quant_config=quant_config, |
| 176 | prefix=f"{prefix}.fc2", |
| 177 | ) |
| 178 | |
| 179 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 180 | x, _ = self.fc1(x) |
| 181 | x = self.act(x) |
| 182 | x, _ = self.fc2(x) |
| 183 | return x |
| 184 | |
| 185 | |
| 186 | class DashengAttention(nn.Module): |
| 187 | def __init__( |
| 188 | self, |
| 189 | dim: int, |
| 190 | num_heads: int = 8, |
| 191 | qkv_bias: bool = False, |
| 192 | quant_config: QuantizationConfig | None = None, |
| 193 | prefix: str = "", |
| 194 | ): |
| 195 | super().__init__() |
| 196 | assert dim % num_heads == 0, "dim should be divisible by num_heads" |
| 197 | self.embed_dim = dim |
| 198 | tp_size = get_tensor_model_parallel_world_size() |
| 199 | self.total_num_heads = num_heads |
| 200 | assert self.total_num_heads % tp_size == 0 |
| 201 | self.num_heads = self.total_num_heads // tp_size |
| 202 | if self.total_num_heads >= tp_size: |
| 203 | # Number of heads is greater than TP size, so we partition |
| 204 | # the KV heads across multiple tensor parallel GPUs. |
| 205 | assert self.total_num_heads % tp_size == 0 |
| 206 | else: |
| 207 | # Number of heads is less than TP size, so we replicate |
| 208 | # the KV heads across multiple tensor parallel GPUs. |
| 209 | assert tp_size % self.total_num_heads == 0 |
| 210 | self.num_kv_heads = max(1, self.total_num_heads // tp_size) |
| 211 | self.head_dim = self.embed_dim // self.total_num_heads |
| 212 | self.q_size = self.num_heads * self.head_dim |
| 213 | self.kv_size = self.num_kv_heads * self.head_dim |
| 214 | self.scale = self.head_dim**-0.5 |
| 215 | |
| 216 | self.qkv = QKVParallelLinear( |
| 217 | hidden_size=self.embed_dim, |
| 218 | head_size=self.head_dim, |
| 219 | total_num_heads=self.total_num_heads, |
| 220 | total_num_kv_heads=self.total_num_heads, |
| 221 | bias=qkv_bias, |
| 222 | quant_config=quant_config, |
| 223 | prefix=f"{prefix}.qkv", |
| 224 | ) |
| 225 | self.proj = RowParallelLinear( |
| 226 | input_size=dim, |
| 227 | output_size=dim, |
| 228 | quant_config=quant_config, |
| 229 | prefix=f"{prefix}.proj", |
| 230 | ) |
| 231 | |
| 232 | def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): |
| 233 | B, N, _ = x.shape |
| 234 | |
| 235 | qkv, _ = self.qkv(x) |
| 236 | qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim) |
| 237 | qkv = qkv.permute(2, 0, 3, 1, 4) |
| 238 | q, k, v = qkv.unbind(0) |
| 239 | |
| 240 | x = scaled_dot_product_attention( |
| 241 | q, |
| 242 | k, |
| 243 | v, |
| 244 | attn_mask=mask[:, None, None, :] if mask is not None else None, |
| 245 | ) |
| 246 | |
| 247 | x = x.transpose(1, 2).reshape(B, N, self.q_size) |
| 248 | x, _ = self.proj(x) |
| 249 | return x |
| 250 | |
| 251 | |
| 252 | class DashengBlock(nn.Module): |
| 253 | def __init__( |
| 254 | self, |
| 255 | dim: int, |
| 256 | num_heads: int, |
| 257 | mlp_ratio: float = 4.0, |
| 258 | qkv_bias: bool = False, |
| 259 | init_values: float | None = None, |
| 260 | quant_config: QuantizationConfig | None = None, |
| 261 | prefix: str = "", |
| 262 | ): |
| 263 | super().__init__() |
| 264 | self.norm1 = nn.LayerNorm(dim, eps=1e-6) |
| 265 | self.attn = DashengAttention( |
| 266 | dim, |
| 267 | num_heads=num_heads, |
| 268 | qkv_bias=qkv_bias, |
| 269 | quant_config=quant_config, |
| 270 | prefix=f"{prefix}.attn", |
| 271 | ) |
| 272 | self.ls1 = ( |
| 273 | LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 274 | ) |
| 275 | |
| 276 | self.norm2 = nn.LayerNorm(dim, eps=1e-6) |
| 277 | self.mlp = DashengMlp( |
| 278 | in_features=dim, |
| 279 | hidden_features=int(dim * mlp_ratio), |
| 280 | quant_config=quant_config, |
| 281 | prefix=f"{prefix}.mlp", |
| 282 | ) |
| 283 | self.ls2 = ( |
| 284 | LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 285 | ) |
| 286 | |
| 287 | # Kwargs usually has a mask parameter that is passed to Attention |
| 288 | def forward( |
| 289 | self, |
| 290 | x: torch.Tensor, |
| 291 | mask: torch.Tensor | None = None, |
| 292 | ) -> torch.Tensor: |
| 293 | x = x + self.ls1(self.attn(self.norm1(x), mask)) |
| 294 | x = x + self.ls2(self.mlp(self.norm2(x))) |
| 295 | return x |
| 296 | |
| 297 | |
| 298 | class DashengFrontend(nn.Module): |
| 299 | def __init__(self, config: DashengConfig): |
| 300 | super().__init__() |
| 301 | self.config = config |
| 302 | |
| 303 | spectrogram_window = torch.hann_window(self.config.win_length) |
| 304 | self.register_buffer( |
| 305 | "spectrogram_window", |
| 306 | spectrogram_window, |
| 307 | persistent=False, |
| 308 | ) |
| 309 | self.spectrogram_window: torch.Tensor |
| 310 | |
| 311 | melscale_fbanks = F.melscale_fbanks( |
| 312 | n_freqs=self.config.n_fft // 2 + 1, |
| 313 | f_min=self.config.f_min, |
| 314 | f_max=self.config.f_max, |
| 315 | n_mels=self.config.n_mels, |
| 316 | sample_rate=self.config.sample_rate, |
| 317 | ) |
| 318 | self.register_buffer("melscale_fbanks", melscale_fbanks, persistent=False) |
| 319 | self.melscale_fbanks: torch.Tensor |
| 320 | |
| 321 | def forward(self, waveform: torch.Tensor) -> torch.Tensor: |
| 322 | spectrogram = F.spectrogram( |
| 323 | waveform=waveform.to(torch.float32), |
| 324 | pad=0, |
| 325 | window=self.spectrogram_window, |
| 326 | n_fft=self.config.n_fft, |
| 327 | hop_length=self.config.hop_length, |
| 328 | win_length=self.config.win_length, |
| 329 | power=2, |
| 330 | normalized=False, |
| 331 | center=self.config.center, |
| 332 | ) |
| 333 | mel_spectrogram = (spectrogram.mT @ self.melscale_fbanks.to(torch.float32)).mT |
| 334 | # x has shape [batch, freq, time]. |
| 335 | # F.amplitude_to_DB accepts inputs shaped as: |
| 336 | # - [freq, time] |
| 337 | # - [channel, freq, time] |
| 338 | # - [..., channel, freq, time] |
| 339 | # Here we insert a channel dimension of size 1 before calling it, |
| 340 | # then remove that extra dimension afterward. |
| 341 | log_mel_spectrogram = F.amplitude_to_DB( |
| 342 | mel_spectrogram.unsqueeze(1), |
| 343 | multiplier=10, |
| 344 | amin=1e-10, |
| 345 | db_multiplier=0, |
| 346 | top_db=120, |
| 347 | ).squeeze(1) |
| 348 | return log_mel_spectrogram.to(waveform.dtype) |
| 349 | |
| 350 | |
| 351 | class DashengAudioTransformer(nn.Module): |
| 352 | def __init__( |
| 353 | self, |
| 354 | config: DashengConfig, |
| 355 | quant_config: QuantizationConfig | None = None, |
| 356 | prefix: str = "", |
| 357 | ): |
| 358 | super().__init__() |
| 359 | |
| 360 | self.target_length = config.target_length |
| 361 | self.hop_length = config.hop_length |
| 362 | |
| 363 | self.front_end = DashengFrontend(config) |
| 364 | |
| 365 | self.init_bn = nn.BatchNorm2d(config.n_mels, momentum=0.01) |
| 366 | |
| 367 | self.patch_embed = AudioPatchEmbed( |
| 368 | input_size=(config.n_mels, config.target_length), |
| 369 | embed_dim=config.embed_dim, |
| 370 | in_chans=config.input_channels, |
| 371 | patch_size=config.patch_size, |
| 372 | flatten=False, |
| 373 | patch_stride=config.patch_stride, |
| 374 | ) |
| 375 | |
| 376 | self.time_pos_embed = nn.Parameter( |
| 377 | torch.empty(1, config.embed_dim, 1, self.patch_embed.grid_size[1]) |
| 378 | ) |
| 379 | self.freq_pos_embed = nn.Parameter( |
| 380 | torch.empty(1, config.embed_dim, self.patch_embed.grid_size[0], 1) |
| 381 | ) |
| 382 | self.blocks = nn.ModuleList( |
| 383 | DashengBlock( |
| 384 | dim=config.embed_dim, |
| 385 | num_heads=config.num_heads, |
| 386 | mlp_ratio=config.mlp_ratio, |
| 387 | qkv_bias=config.qkv_bias, |
| 388 | init_values=config.init_values, |
| 389 | quant_config=quant_config, |
| 390 | prefix=f"{prefix}.blocks.{i}", |
| 391 | ) |
| 392 | for i in range(config.depth) |
| 393 | ) |
| 394 | self.norm = nn.LayerNorm(config.embed_dim, eps=1e-6) |
| 395 | |
| 396 | def forward_features( |
| 397 | self, |
| 398 | x: torch.Tensor, |
| 399 | mask: torch.Tensor | None = None, |
| 400 | ) -> torch.Tensor: |
| 401 | t = x.shape[-1] |
| 402 | x = x + self.time_pos_embed[:, :, :, :t] |
| 403 | x = ( |
| 404 | x + self.freq_pos_embed[:, :, :, :] |
| 405 | ) # Just to support __getitem__ in posembed |
| 406 | x = torch.permute( |
| 407 | torch.flatten(x, 2, 3), (0, 2, 1) |
| 408 | ) # rearrange(x, "b c f t -> b (f t) c") |
| 409 | for block in self.blocks: |
| 410 | x = block(x, mask) |
| 411 | x = self.norm(x) |
| 412 | return x |
| 413 | |
| 414 | def _to_mask(self, lengths: torch.Tensor, max_length: int) -> torch.Tensor: |
| 415 | batch_size = len(lengths) |
| 416 | idx = torch.arange(max_length, device=lengths.device) |
| 417 | idx = idx.repeat(batch_size).view(batch_size, max_length) |
| 418 | mask = (idx < lengths.unsqueeze(-1)).bool() |
| 419 | return mask |
| 420 | |
| 421 | def forward( |
| 422 | self, |
| 423 | x: torch.Tensor, |
| 424 | x_length: torch.Tensor | None = None, |
| 425 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 426 | x = self.front_end(x) |
| 427 | x = x.to(self.time_pos_embed.dtype) |
| 428 | target_length_in_patches = self.target_length // 4 |
| 429 | x = x.unsqueeze(1) |
| 430 | x = torch.permute(x, (0, 2, 1, 3)) |
| 431 | x = self.init_bn(x) |
| 432 | x = torch.permute(x, (0, 2, 1, 3)) |
| 433 | |
| 434 | x = self.patch_embed(x) |
| 435 | t = x.shape[-1] |
| 436 | |
| 437 | input_splits = x.split(target_length_in_patches, dim=-1) |
| 438 | |
| 439 | if x_length is not None: |
| 440 | assert len(x_length) == len(x), ( |
| 441 | "batchsizes of input x and x_length need to be same" |
| 442 | ) |
| 443 | assert x_length.ndim == 1, "Lengths are of size (B,)" |
| 444 | scaled_lengths = (x_length / (self.hop_length * 4)).long() |
| 445 | mask = self._to_mask(max_length=t, lengths=scaled_lengths) |
| 446 | split_masks = mask.split(target_length_in_patches, dim=-1) |
| 447 | else: |
| 448 | mask = None |
| 449 | split_masks = [None] * len(input_splits) |
| 450 | |
| 451 | outputs = [] |
| 452 | |
| 453 | for split_x, split_mask in zip(input_splits, split_masks): |
| 454 | forward_kwargs = {} |
| 455 | forward_kwargs["mask"] = split_mask |
| 456 | split_x = self.forward_features(split_x, **forward_kwargs) |
| 457 | outputs.append(split_x) |
| 458 | x = torch.cat(outputs, dim=1) |
| 459 | return x, mask |
| 460 | |
| 461 | |
| 462 | class AudioProjectorSubsample(nn.Module): |
| 463 | def __init__( |
| 464 | self, |
| 465 | in_dim: int, |
| 466 | out_dim: int, |
| 467 | downsample_rate=5, |
| 468 | dtype: torch.dtype | None = None, |
| 469 | quant_config: QuantizationConfig | None = None, |
| 470 | prefix: str = "", |
| 471 | ): |
| 472 | super().__init__() |
| 473 | self.k = downsample_rate |
| 474 | self.net = nn.Sequential( |
| 475 | ColumnParallelLinear( |
| 476 | input_size=in_dim * self.k, |
| 477 | output_size=out_dim, |
| 478 | quant_config=quant_config, |
| 479 | prefix=f"{prefix}.net.0", |
| 480 | return_bias=False, |
| 481 | ), |
| 482 | get_act_fn("gelu"), |
| 483 | RowParallelLinear( |
| 484 | input_size=out_dim, |
| 485 | output_size=out_dim, |
| 486 | quant_config=quant_config, |
| 487 | prefix=f"{prefix}.net.2", |
| 488 | return_bias=False, |
| 489 | ), |
| 490 | ) |
| 491 | |
| 492 | def forward(self, x, mask=None): |
| 493 | batch_size, seq_len, dim = x.shape |
| 494 | num_frames_to_discard = seq_len % self.k |
| 495 | if num_frames_to_discard > 0: |
| 496 | x = x[:, :-num_frames_to_discard, :] |
| 497 | if mask is not None: |
| 498 | mask = mask[:, :-num_frames_to_discard] |
| 499 | if mask is None: |
| 500 | mask = torch.ones(x.shape[:-1], dtype=torch.long, device=x.device) |
| 501 | x = x.reshape( |
| 502 | batch_size, -1, self.k * dim |
| 503 | ) # rearrange(x, "b (s k) d -> b s (k d)", k=self.k) |
| 504 | for layer in self.net: |
| 505 | x = layer(x) |
| 506 | mask = mask.reshape( |
| 507 | batch_size, -1, self.k |
| 508 | ) # rearrange(mask, "b (s k) -> b s k", k=self.k) |
| 509 | mask = mask.any(dim=-1).long() |
| 510 | return x, mask |
| 511 | |
| 512 | |
| 513 | # === Audio Inputs === # |
| 514 | class MiDashengLMAudioInputs(TensorSchema): |
| 515 | """Dimensions: |
| 516 | - bn: Batch size * number of audios |
| 517 | - p: Number of sampling points |
| 518 | """ |
| 519 | |
| 520 | input_values: Annotated[torch.Tensor, TensorShape("n", "p")] |
| 521 | audio_length: Annotated[torch.Tensor, TensorShape("n")] |
| 522 | |
| 523 | |
| 524 | class MiDashengLMProcessingInfo(BaseProcessingInfo): |
| 525 | def get_hf_config(self): |
| 526 | return self.ctx.get_hf_config() |
| 527 | |
| 528 | def get_feature_extractor(self): |
| 529 | hf_processor = self.get_hf_processor() |
| 530 | feature_extractor = hf_processor.feature_extractor |
| 531 | return feature_extractor |
| 532 | |
| 533 | def get_data_parser(self): |
| 534 | feature_extractor = self.get_feature_extractor() |
| 535 | |
| 536 | return MultiModalDataParser( |
| 537 | target_sr=feature_extractor.sampling_rate, |
| 538 | expected_hidden_size=self._get_expected_hidden_size(), |
| 539 | ) |
| 540 | |
| 541 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 542 | return {"audio": None} |
| 543 | |
| 544 | def get_min_audio_len(self): |
| 545 | return 3200 |
| 546 | |
| 547 | def get_max_audio_len(self): |
| 548 | return 160000 |
| 549 | |
| 550 | |
| 551 | class MiDashengLMDummyInputsBuilder(BaseDummyInputsBuilder[MiDashengLMProcessingInfo]): |
| 552 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 553 | num_audios = mm_counts.get("audio", 0) |
| 554 | |
| 555 | hf_processor = self.info.get_hf_processor() |
| 556 | audio_token = hf_processor.audio_token |
| 557 | audio_bos_token = hf_processor.audio_bos_token |
| 558 | audio_eos_token = hf_processor.audio_eos_token |
| 559 | |
| 560 | single_audio_text = f"{audio_bos_token}{audio_token}{audio_eos_token}" |
| 561 | return single_audio_text * num_audios |
| 562 | |
| 563 | def get_dummy_mm_data( |
| 564 | self, |
| 565 | seq_len: int, |
| 566 | mm_counts: Mapping[str, int], |
| 567 | mm_options: MultiModalDummyOptions, |
| 568 | ) -> MultiModalDataDict: |
| 569 | return { |
| 570 | "audio": self._get_dummy_audios( |
| 571 | length=self.info.get_max_audio_len(), |
| 572 | num_audios=mm_counts.get("audio", 0), |
| 573 | overrides=mm_options.get("audio"), |
| 574 | ) |
| 575 | } |
| 576 | |
| 577 | |
| 578 | class MiDashengLMMultiModalProcessor( |
| 579 | BaseMultiModalProcessor[MiDashengLMProcessingInfo] |
| 580 | ): |
| 581 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 582 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 583 | |
| 584 | def _get_hf_mm_inputs( |
| 585 | self, |
| 586 | mm_items: MultiModalDataItems, |
| 587 | hf_kwargs: Mapping[str, object], |
| 588 | ) -> HFMultiModalInputs: |
| 589 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 590 | |
| 591 | if audios := hf_inputs.hf_data.get("audio"): |
| 592 | assert isinstance(audios, list) |
| 593 | |
| 594 | min_audio_len = self.info.get_min_audio_len() |
| 595 | hf_inputs.hf_data["audio"] = [ |
| 596 | np.pad( |
| 597 | audio, |
| 598 | (0, min_audio_len - audio.shape[-1]), |
| 599 | mode="constant", |
| 600 | constant_values=0, |
| 601 | ) |
| 602 | if isinstance(audio, np.ndarray) and audio.shape[-1] < min_audio_len |
| 603 | else audio |
| 604 | for audio in audios |
| 605 | ] |
| 606 | |
| 607 | return hf_inputs |
| 608 | |
| 609 | def _get_mm_fields_config( |
| 610 | self, |
| 611 | hf_inputs: BatchFeature, |
| 612 | hf_processor_mm_kwargs: Mapping[str, object], |
| 613 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 614 | return dict( |
| 615 | input_values=MultiModalFieldConfig.batched("audio"), |
| 616 | audio_length=MultiModalFieldConfig.batched("audio"), |
| 617 | ) |
| 618 | |
| 619 | def _get_prompt_updates( |
| 620 | self, |
| 621 | mm_items: MultiModalDataItems, |
| 622 | hf_processor_mm_kwargs: Mapping[str, object], |
| 623 | out_mm_kwargs: MultiModalKwargsItems, |
| 624 | ) -> Sequence[PromptUpdate]: |
| 625 | processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 626 | tokenizer = self.info.get_tokenizer() |
| 627 | vocab = tokenizer.get_vocab() |
| 628 | |
| 629 | audio_token = getattr(processor, "audio_token", "<|AUDIO|>") |
| 630 | audio_token_id = vocab[audio_token] |
| 631 | |
| 632 | out_mm_data = out_mm_kwargs.get_data() |
| 633 | audio_length = out_mm_data.get("audio_length") |
| 634 | audio_output_lengths: list[int] |
| 635 | if audio_length is None: |
| 636 | audio_output_lengths = [] |
| 637 | else: |
| 638 | audio_length_np = ( |
| 639 | audio_length.cpu().numpy() |
| 640 | if isinstance(audio_length, torch.Tensor) |
| 641 | else audio_length |
| 642 | ) |
| 643 | assert isinstance(audio_length_np, (list, tuple, np.ndarray)) |
| 644 | audio_output_lengths = [] |
| 645 | for length in audio_length_np: |
| 646 | assert isinstance(length, (int, np.integer)) |
| 647 | audio_output_lengths.append( |
| 648 | max(1, calculate_mel_frames_dasheng(int(length))) |
| 649 | ) |
| 650 | |
| 651 | def get_replacement_midashenglm(item_idx: int): |
| 652 | num_features = audio_output_lengths[item_idx] |
| 653 | audio_tokens = [audio_token_id] * num_features |
| 654 | |
| 655 | return PromptUpdateDetails.select_token_id( |
| 656 | audio_tokens, |
| 657 | embed_token_id=audio_token_id, |
| 658 | ) |
| 659 | |
| 660 | return [ |
| 661 | PromptReplacement( |
| 662 | modality="audio", |
| 663 | target=[audio_token_id], |
| 664 | replacement=get_replacement_midashenglm, |
| 665 | ) |
| 666 | ] |
| 667 | |
| 668 | |
| 669 | @MULTIMODAL_REGISTRY.register_processor( |
| 670 | MiDashengLMMultiModalProcessor, |
| 671 | info=MiDashengLMProcessingInfo, |
| 672 | dummy_inputs=MiDashengLMDummyInputsBuilder, |
| 673 | ) |
| 674 | class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP): |
| 675 | packed_modules_mapping = { |
| 676 | "qkv_proj": [ |
| 677 | "q_proj", |
| 678 | "k_proj", |
| 679 | "v_proj", |
| 680 | ], |
| 681 | "gate_up_proj": [ |
| 682 | "gate_proj", |
| 683 | "up_proj", |
| 684 | ], |
| 685 | } |
| 686 | |
| 687 | @classmethod |
| 688 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 689 | if modality.startswith("audio"): |
| 690 | return "<|audio_bos|><|AUDIO|><|audio_eos|>" |
| 691 | |
| 692 | raise ValueError("Only audio modality is supported") |
| 693 | |
| 694 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 695 | super().__init__() |
| 696 | config = vllm_config.model_config.hf_config |
| 697 | quant_config = vllm_config.quant_config |
| 698 | self.config = config |
| 699 | self.quant_config = quant_config |
| 700 | |
| 701 | with self._mark_tower_model(vllm_config, "audio"): |
| 702 | self.audio_encoder = DashengAudioTransformer( |
| 703 | config.audio_encoder_config, |
| 704 | quant_config=quant_config, |
| 705 | prefix=maybe_prefix(prefix, "audio_encoder"), |
| 706 | ) |
| 707 | self.audio_projector = AudioProjectorSubsample( |
| 708 | in_dim=config.audio_encoder_config.embed_dim, |
| 709 | out_dim=config.text_config.hidden_size, |
| 710 | downsample_rate=config.subsample_factor, |
| 711 | quant_config=quant_config, |
| 712 | prefix=maybe_prefix(prefix, "audio_projector"), |
| 713 | ) |
| 714 | |
| 715 | with self._mark_language_model(vllm_config): |
| 716 | self.decoder = init_vllm_registered_model( |
| 717 | vllm_config=vllm_config, |
| 718 | hf_config=config.text_config, |
| 719 | prefix=maybe_prefix(prefix, "decoder"), |
| 720 | architectures=["Qwen2ForCausalLM"], |
| 721 | ) |
| 722 | |
| 723 | self.make_empty_intermediate_tensors = ( |
| 724 | self.decoder.make_empty_intermediate_tensors |
| 725 | ) |
| 726 | |
| 727 | def _parse_and_validate_audio_input( |
| 728 | self, **kwargs: object |
| 729 | ) -> MiDashengLMAudioInputs | None: |
| 730 | input_values = kwargs.pop("input_values", None) |
| 731 | audio_length = kwargs.pop("audio_length", None) |
| 732 | |
| 733 | if input_values is None: |
| 734 | return None |
| 735 | |
| 736 | if isinstance(input_values, list): |
| 737 | input_values = torch.nn.utils.rnn.pad_sequence( |
| 738 | input_values, |
| 739 | batch_first=True, |
| 740 | ) |
| 741 | |
| 742 | return MiDashengLMAudioInputs( |
| 743 | input_values=input_values, |
| 744 | audio_length=audio_length, |
| 745 | ) |
| 746 | |
| 747 | def _process_audio_input( |
| 748 | self, |
| 749 | audio_input: MiDashengLMAudioInputs, |
| 750 | ) -> tuple[torch.Tensor, ...]: |
| 751 | # Process audio through encoder and projector |
| 752 | input_values = audio_input["input_values"] |
| 753 | audio_length = audio_input["audio_length"] |
| 754 | |
| 755 | encoder_out, encoder_atts = self.audio_encoder(input_values, audio_length) |
| 756 | audio_embeddings, _ = self.audio_projector(encoder_out, encoder_atts) |
| 757 | audio_embeddings = audio_embeddings.to(audio_input["input_values"].dtype) |
| 758 | batch_size, max_audio_tokens, embed_dim = audio_embeddings.shape |
| 759 | |
| 760 | audio_output_lengths = [ |
| 761 | max(1, calculate_mel_frames_dasheng(int(length))) # at least one frame |
| 762 | for length in audio_length.tolist() |
| 763 | ] |
| 764 | audio_output_lengths = torch.tensor( |
| 765 | audio_output_lengths, |
| 766 | device=audio_embeddings.device, |
| 767 | ) |
| 768 | |
| 769 | audio_feature_mask = torch.arange( |
| 770 | max_audio_tokens, device=audio_embeddings.device |
| 771 | ).unsqueeze(0).expand( |
| 772 | batch_size, max_audio_tokens |
| 773 | ) < audio_output_lengths.unsqueeze(1) |
| 774 | |
| 775 | masked_audio_features = audio_embeddings[audio_feature_mask].view(-1, embed_dim) |
| 776 | |
| 777 | return torch.split(masked_audio_features, audio_output_lengths.tolist()) |
| 778 | |
| 779 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 780 | audio_input = self._parse_and_validate_audio_input(**kwargs) |
| 781 | |
| 782 | if audio_input is None: |
| 783 | return [] |
| 784 | return self._process_audio_input(audio_input) |
| 785 | |
| 786 | def forward( |
| 787 | self, |
| 788 | input_ids: torch.Tensor | None, |
| 789 | positions: torch.Tensor, |
| 790 | intermediate_tensors: IntermediateTensors | None = None, |
| 791 | inputs_embeds: torch.Tensor | None = None, |
| 792 | **kwargs: object, |
| 793 | ) -> torch.Tensor | IntermediateTensors: |
| 794 | if intermediate_tensors is not None: |
| 795 | inputs_embeds = None |
| 796 | |
| 797 | return self.decoder.model( |
| 798 | input_ids, |
| 799 | positions, |
| 800 | intermediate_tensors, |
| 801 | inputs_embeds=inputs_embeds, |
| 802 | ) |
| 803 | |
| 804 | def compute_logits( |
| 805 | self, |
| 806 | hidden_states: torch.Tensor, |
| 807 | ) -> torch.Tensor | None: |
| 808 | return self.decoder.compute_logits(hidden_states) |
| 809 | |
| 810 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 811 | loader = AutoWeightsLoader(self) |
| 812 | return loader.load_weights(weights) |