| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Inference-only MOSS-Transcribe-Diarize ASR model. |
| 5 | |
| 6 | The checkpoint layout is: |
| 7 | |
| 8 | * ``model.whisper_encoder.*``: Whisper-medium encoder weights. |
| 9 | * ``model.vq_adaptor.*``: 4x time-merge projector. |
| 10 | * ``model.language_model.*``: Qwen3-0.6B decoder weights. |
| 11 | """ |
| 12 | |
| 13 | import math |
| 14 | from collections.abc import Iterable, Mapping, Sequence |
| 15 | from typing import Annotated, Any, Literal, TypeAlias |
| 16 | |
| 17 | import regex as re |
| 18 | import torch |
| 19 | from torch import nn |
| 20 | from transformers import BatchFeature |
| 21 | |
| 22 | from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig |
| 23 | from vllm.config.multimodal import MultiModalDummyOptions |
| 24 | from vllm.config.speech_to_text import SpeechToTextParams |
| 25 | from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt |
| 26 | from vllm.model_executor.models.interfaces import ( |
| 27 | DiarizedTranscriptionSegment, |
| 28 | MultiModalEmbeddings, |
| 29 | SupportsMultiModal, |
| 30 | SupportsPP, |
| 31 | SupportsTranscription, |
| 32 | _require_is_multimodal, |
| 33 | ) |
| 34 | from vllm.model_executor.models.utils import ( |
| 35 | AutoWeightsLoader, |
| 36 | WeightsMapper, |
| 37 | _merge_multimodal_embeddings, |
| 38 | init_vllm_registered_model, |
| 39 | maybe_prefix, |
| 40 | parse_diarized_speaker, |
| 41 | parse_diarized_timestamp, |
| 42 | ) |
| 43 | from vllm.model_executor.models.whisper import ( |
| 44 | WhisperEncoder, |
| 45 | _create_fake_bias_for_k_proj, |
| 46 | ) |
| 47 | from vllm.model_executor.models.whisper_utils import ISO639_1_SUPPORTED_LANGS |
| 48 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 49 | from vllm.multimodal.inputs import ( |
| 50 | AudioItem, |
| 51 | MultiModalFieldConfig, |
| 52 | MultiModalKwargsItems, |
| 53 | ) |
| 54 | from vllm.multimodal.parse import ( |
| 55 | DictEmbeddingItems, |
| 56 | ModalityDataItems, |
| 57 | MultiModalDataItems, |
| 58 | MultiModalDataParser, |
| 59 | ) |
| 60 | from vllm.multimodal.processing import ( |
| 61 | BaseDummyInputsBuilder, |
| 62 | BaseMultiModalProcessor, |
| 63 | BaseProcessingInfo, |
| 64 | PromptReplacement, |
| 65 | PromptUpdate, |
| 66 | PromptUpdateDetails, |
| 67 | cached_encode, |
| 68 | ) |
| 69 | from vllm.multimodal.processing.processor import HFMultiModalInputs, ProcessorInputs |
| 70 | from vllm.sequence import IntermediateTensors |
| 71 | from vllm.transformers_utils.processor import cached_processor_from_config |
| 72 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 73 | |
| 74 | WHISPER_ENCODER_STRIDE = 2 |
| 75 | MAX_AUDIO_DURATION_S = 90 * 60 |
| 76 | |
| 77 | AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>" |
| 78 | |
| 79 | DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT = ( |
| 80 | "请将音频转写为文本，每一段需以起始时间戳和说话人编号" |
| 81 | "（[S01]、[S02]、[S03]…）开头，正文为对应的语音内容，" |
| 82 | "并在段末标注结束时间戳，以清晰标明该段语音范围。" |
| 83 | ) |
| 84 | |
| 85 | _MOSS_DIARIZED_HEADER_RE = re.compile( |
| 86 | r"\[(?P<start>[0-9.]{1,32})\]\s*\[(?P<speaker>S[0-9]{1,15})\]" |
| 87 | ) |
| 88 | _MOSS_DIARIZED_END_RE = re.compile(r"\[(?P<end>[0-9.]{1,32})\]\s*\Z") |
| 89 | |
| 90 | |
| 91 | class MossTranscribeDiarizeAudioInputs(TensorSchema): |
| 92 | """Dimensions: |
| 93 | - c: Audio chunks |
| 94 | - m: Mel bins |
| 95 | - f: Mel frames |
| 96 | - n: Number of audio items |
| 97 | """ |
| 98 | |
| 99 | type: Literal["audio_features"] = "audio_features" |
| 100 | |
| 101 | input_features: Annotated[ |
| 102 | torch.Tensor | None, |
| 103 | TensorShape("c", "m", "f"), |
| 104 | ] |
| 105 | audio_feature_lengths: Annotated[ |
| 106 | torch.Tensor | None, |
| 107 | TensorShape("c"), |
| 108 | ] |
| 109 | audio_chunk_counts: Annotated[ |
| 110 | torch.Tensor | None, |
| 111 | TensorShape("n"), |
| 112 | ] |
| 113 | |
| 114 | |
| 115 | class MossTranscribeDiarizeEmbeddingInputs(TensorSchema): |
| 116 | """Dimensions: |
| 117 | - n: Number of audio items |
| 118 | - t: Number of audio tokens |
| 119 | - h: Hidden size |
| 120 | """ |
| 121 | |
| 122 | type: Literal["audio_embeds"] = "audio_embeds" |
| 123 | |
| 124 | audio_embeds: Annotated[ |
| 125 | list[torch.Tensor], |
| 126 | TensorShape("n", "t", "h", dynamic_dims={"t"}), |
| 127 | ] |
| 128 | |
| 129 | |
| 130 | MossTranscribeDiarizeInputs: TypeAlias = ( |
| 131 | MossTranscribeDiarizeAudioInputs | MossTranscribeDiarizeEmbeddingInputs |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | def _compute_total_audio_tokens( |
| 136 | num_samples: int, |
| 137 | feature_extractor: Any, |
| 138 | audio_merge_size: int, |
| 139 | ) -> int: |
| 140 | if num_samples <= 0: |
| 141 | return 0 |
| 142 | |
| 143 | n_samples = int(feature_extractor.n_samples) |
| 144 | stride = ( |
| 145 | int(feature_extractor.hop_length) |
| 146 | * WHISPER_ENCODER_STRIDE |
| 147 | * int(audio_merge_size) |
| 148 | ) |
| 149 | total = 0 |
| 150 | for start in range(0, num_samples, n_samples): |
| 151 | chunk_samples = min(n_samples, num_samples - start) |
| 152 | total += (chunk_samples - 1) // stride + 1 |
| 153 | return total |
| 154 | |
| 155 | |
| 156 | def _get_max_audio_samples(feature_extractor: Any) -> int: |
| 157 | return int(MAX_AUDIO_DURATION_S * feature_extractor.sampling_rate) |
| 158 | |
| 159 | |
| 160 | def _as_audio_embedding_list(audio_embeds: object) -> list[torch.Tensor]: |
| 161 | if isinstance(audio_embeds, torch.Tensor): |
| 162 | if audio_embeds.ndim == 2: |
| 163 | return [audio_embeds] |
| 164 | if audio_embeds.ndim == 3: |
| 165 | return list(audio_embeds.unbind(dim=0)) |
| 166 | raise ValueError( |
| 167 | f"`audio_embeds` must be a 2D or 3D tensor, got {audio_embeds.ndim}D." |
| 168 | ) |
| 169 | |
| 170 | if isinstance(audio_embeds, (list, tuple)) and all( |
| 171 | isinstance(audio_embed, torch.Tensor) for audio_embed in audio_embeds |
| 172 | ): |
| 173 | return list(audio_embeds) |
| 174 | |
| 175 | raise TypeError( |
| 176 | "`audio_embeds` must be a torch.Tensor or a list of torch.Tensor " |
| 177 | f"objects, got {type(audio_embeds)!r}." |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | def _get_required_token_id(tokenizer: Any, token: str) -> int: |
| 182 | token_id = tokenizer.convert_tokens_to_ids(token) |
| 183 | if token_id is None or (isinstance(token_id, int) and token_id < 0): |
| 184 | raise ValueError(f"Tokenizer is missing required token {token!r}.") |
| 185 | return int(token_id) |
| 186 | |
| 187 | |
| 188 | def _get_audios_from_mm_data(mm_data: Mapping[str, object]) -> list[Any]: |
| 189 | mm_data_dict = dict(mm_data) |
| 190 | audio_data = mm_data_dict.pop("audios", mm_data_dict.pop("audio", [])) |
| 191 | if isinstance(audio_data, list): |
| 192 | audios = audio_data |
| 193 | elif isinstance(audio_data, tuple) and not ( |
| 194 | len(audio_data) == 2 and isinstance(audio_data[1], (int, float)) |
| 195 | ): |
| 196 | audios = list(audio_data) |
| 197 | elif audio_data is None: |
| 198 | audios = [] |
| 199 | else: |
| 200 | audios = [audio_data] |
| 201 | |
| 202 | audio_arrays: list[Any] = [] |
| 203 | for audio in audios: |
| 204 | if isinstance(audio, (tuple, list)) and len(audio) == 2: |
| 205 | audio = audio[0] |
| 206 | audio_arrays.append(audio) |
| 207 | return audio_arrays |
| 208 | |
| 209 | |
| 210 | def _add_vllm_audio_metadata( |
| 211 | processed: BatchFeature, |
| 212 | num_audios: int, |
| 213 | ) -> BatchFeature: |
| 214 | audio_feature_lengths = processed["audio_feature_lengths"] |
| 215 | if not isinstance(audio_feature_lengths, torch.Tensor): |
| 216 | audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long) |
| 217 | audio_feature_lengths = audio_feature_lengths.to(dtype=torch.long) |
| 218 | |
| 219 | audio_chunk_mapping = processed.get("audio_chunk_mapping") |
| 220 | if audio_chunk_mapping is None: |
| 221 | if num_audios == 1: |
| 222 | audio_chunk_mapping = torch.zeros_like(audio_feature_lengths) |
| 223 | elif audio_feature_lengths.numel() == num_audios: |
| 224 | audio_chunk_mapping = torch.arange(num_audios, dtype=torch.long) |
| 225 | else: |
| 226 | raise ValueError( |
| 227 | "The MOSS processor did not return `audio_chunk_mapping`, and " |
| 228 | "the chunk-to-audio mapping cannot be inferred." |
| 229 | ) |
| 230 | elif not isinstance(audio_chunk_mapping, torch.Tensor): |
| 231 | audio_chunk_mapping = torch.tensor(audio_chunk_mapping, dtype=torch.long) |
| 232 | audio_chunk_mapping = audio_chunk_mapping.to(dtype=torch.long) |
| 233 | if audio_chunk_mapping.numel() != audio_feature_lengths.numel(): |
| 234 | raise ValueError( |
| 235 | "`audio_chunk_mapping` must contain one item per audio chunk: got " |
| 236 | f"{audio_chunk_mapping.numel()} mappings for " |
| 237 | f"{audio_feature_lengths.numel()} chunks." |
| 238 | ) |
| 239 | if audio_chunk_mapping.numel() > 0 and ( |
| 240 | audio_chunk_mapping.min().item() < 0 |
| 241 | or audio_chunk_mapping.max().item() >= num_audios |
| 242 | ): |
| 243 | raise ValueError("`audio_chunk_mapping` contains an out-of-range audio index.") |
| 244 | |
| 245 | audio_chunk_counts = torch.bincount( |
| 246 | audio_chunk_mapping.cpu(), |
| 247 | minlength=num_audios, |
| 248 | ).to(dtype=torch.long) |
| 249 | audio_token_lengths = torch.zeros( |
| 250 | num_audios, |
| 251 | dtype=torch.long, |
| 252 | device=audio_feature_lengths.device, |
| 253 | ) |
| 254 | audio_token_lengths.scatter_add_( |
| 255 | 0, |
| 256 | audio_chunk_mapping.to(device=audio_feature_lengths.device), |
| 257 | audio_feature_lengths, |
| 258 | ) |
| 259 | |
| 260 | processed["audio_feature_lengths"] = audio_feature_lengths |
| 261 | processed["audio_chunk_counts"] = audio_chunk_counts |
| 262 | processed["audio_token_lengths"] = audio_token_lengths |
| 263 | return processed |
| 264 | |
| 265 | |
| 266 | class MossTranscribeDiarizeWhisperEncoder(WhisperEncoder): |
| 267 | packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} |
| 268 | hf_to_vllm_mapper = WeightsMapper( |
| 269 | orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}, |
| 270 | orig_to_new_stacked={ |
| 271 | ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), |
| 272 | ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), |
| 273 | ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), |
| 274 | }, |
| 275 | ) |
| 276 | |
| 277 | def __init__( |
| 278 | self, |
| 279 | *, |
| 280 | vllm_config: VllmConfig, |
| 281 | prefix: str = "", |
| 282 | ) -> None: |
| 283 | super().__init__( |
| 284 | vllm_config=vllm_config.with_hf_config( |
| 285 | vllm_config.model_config.hf_config.audio_config |
| 286 | ), |
| 287 | prefix=prefix, |
| 288 | ) |
| 289 | self.audio_merge_size = int(vllm_config.model_config.hf_config.audio_merge_size) |
| 290 | self.max_encoder_batch: int | None = None |
| 291 | |
| 292 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 293 | weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight") |
| 294 | loader = AutoWeightsLoader(self) |
| 295 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 296 | |
| 297 | def forward( |
| 298 | self, |
| 299 | input_features: torch.Tensor, |
| 300 | audio_feature_lengths: torch.Tensor | None = None, |
| 301 | ) -> torch.Tensor: |
| 302 | assert audio_feature_lengths is not None |
| 303 | if input_features.numel() == 0: |
| 304 | return input_features.new_empty((1, 0, self.conv1.out_channels)) |
| 305 | device = self.conv1.weight.device |
| 306 | dtype = self.conv1.weight.dtype |
| 307 | input_features = input_features.to(device=device, dtype=dtype) |
| 308 | audio_feature_lengths = audio_feature_lengths.to(device=device) |
| 309 | batch_size = self.max_encoder_batch or input_features.shape[0] |
| 310 | encoded_parts: list[torch.Tensor] = [] |
| 311 | |
| 312 | for start in range(0, input_features.shape[0], batch_size): |
| 313 | encoded_parts.append( |
| 314 | super().forward([input_features[start : start + batch_size]]) |
| 315 | ) |
| 316 | |
| 317 | hidden = torch.cat(encoded_parts, dim=0) |
| 318 | chunks = [ |
| 319 | hidden[idx : idx + 1, : int(token_len.item()) * self.audio_merge_size] |
| 320 | for idx, token_len in enumerate(audio_feature_lengths) |
| 321 | ] |
| 322 | return torch.cat(chunks, dim=1) |
| 323 | |
| 324 | |
| 325 | class MossTranscribeDiarizeVQAdaptor(nn.Module): |
| 326 | def __init__(self, input_dim: int, hidden_size: int, eps: float) -> None: |
| 327 | super().__init__() |
| 328 | self.layers = nn.Sequential( |
| 329 | nn.Linear(input_dim, hidden_size, bias=True), |
| 330 | nn.SiLU(), |
| 331 | nn.Linear(hidden_size, hidden_size, bias=True), |
| 332 | nn.LayerNorm(hidden_size, eps=eps, bias=True), |
| 333 | ) |
| 334 | |
| 335 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 336 | return self.layers(x) |
| 337 | |
| 338 | |
| 339 | def _mtd_field_config( |
| 340 | hf_inputs: Mapping[str, torch.Tensor], |
| 341 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 342 | fields: dict[str, MultiModalFieldConfig] = {} |
| 343 | if "audio_embeds" in hf_inputs: |
| 344 | fields["audio_embeds"] = MultiModalFieldConfig.batched("audio") |
| 345 | if "audio_chunk_counts" in hf_inputs: |
| 346 | audio_chunk_counts = hf_inputs["audio_chunk_counts"] |
| 347 | fields.update( |
| 348 | input_features=MultiModalFieldConfig.flat_from_sizes( |
| 349 | "audio", |
| 350 | audio_chunk_counts, |
| 351 | ), |
| 352 | audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes( |
| 353 | "audio", |
| 354 | audio_chunk_counts, |
| 355 | ), |
| 356 | audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), |
| 357 | audio_token_lengths=MultiModalFieldConfig.batched( |
| 358 | "audio", keep_on_cpu=True |
| 359 | ), |
| 360 | ) |
| 361 | return fields |
| 362 | |
| 363 | |
| 364 | class MossTranscribeDiarizeMultiModalDataParser(MultiModalDataParser): |
| 365 | def _parse_audio_data( |
| 366 | self, |
| 367 | data: dict[str, torch.Tensor] | ModalityData[AudioItem], |
| 368 | ) -> ModalityDataItems[Any, Any] | None: |
| 369 | if isinstance(data, dict): |
| 370 | return DictEmbeddingItems( |
| 371 | data, |
| 372 | modality="audio", |
| 373 | required_fields={"audio_embeds"}, |
| 374 | fields_factory=_mtd_field_config, |
| 375 | ) |
| 376 | |
| 377 | return super()._parse_audio_data(data) |
| 378 | |
| 379 | |
| 380 | class MossTranscribeDiarizeProcessingInfo(BaseProcessingInfo): |
| 381 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 382 | return {"audio": 1} |
| 383 | |
| 384 | def get_hf_processor(self, **kwargs: object) -> Any: |
| 385 | return self.ctx.get_hf_processor(**kwargs) |
| 386 | |
| 387 | def get_feature_extractor(self, **kwargs: object) -> Any: |
| 388 | return self.get_hf_processor(**kwargs).feature_extractor |
| 389 | |
| 390 | def get_data_parser(self) -> MultiModalDataParser: |
| 391 | feature_extractor = self.get_feature_extractor() |
| 392 | return MossTranscribeDiarizeMultiModalDataParser( |
| 393 | target_sr=feature_extractor.sampling_rate, |
| 394 | target_channels=1, |
| 395 | expected_hidden_size=self._get_expected_hidden_size(), |
| 396 | ) |
| 397 | |
| 398 | def get_mm_max_tokens_per_item( |
| 399 | self, |
| 400 | seq_len: int, |
| 401 | mm_counts: Mapping[str, int], |
| 402 | ) -> Mapping[str, int]: |
| 403 | if mm_counts.get("audio", 0) <= 0: |
| 404 | return {} |
| 405 | |
| 406 | feature_extractor = self.get_feature_extractor() |
| 407 | max_audio_samples = _get_max_audio_samples(feature_extractor) |
| 408 | max_audio_tokens = _compute_total_audio_tokens( |
| 409 | max_audio_samples, |
| 410 | feature_extractor, |
| 411 | self.get_hf_processor().audio_merge_size, |
| 412 | ) |
| 413 | return {"audio": min(seq_len, max_audio_tokens)} |
| 414 | |
| 415 | |
| 416 | class MossTranscribeDiarizeDummyInputsBuilder( |
| 417 | BaseDummyInputsBuilder[MossTranscribeDiarizeProcessingInfo] |
| 418 | ): |
| 419 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 420 | return AUDIO_PLACEHOLDER * mm_counts.get("audio", 0) |
| 421 | |
| 422 | def get_dummy_mm_data( |
| 423 | self, |
| 424 | seq_len: int, |
| 425 | mm_counts: Mapping[str, int], |
| 426 | mm_options: MultiModalDummyOptions, |
| 427 | ) -> MultiModalDataDict: |
| 428 | num_audios = mm_counts.get("audio", 0) |
| 429 | if num_audios == 0: |
| 430 | return {} |
| 431 | |
| 432 | feature_extractor = self.info.get_feature_extractor() |
| 433 | return { |
| 434 | "audio": self._get_dummy_audios( |
| 435 | length=_get_max_audio_samples(feature_extractor), |
| 436 | num_audios=num_audios, |
| 437 | overrides=mm_options.get("audio"), |
| 438 | ) |
| 439 | } |
| 440 | |
| 441 | |
| 442 | class MossTranscribeDiarizeMultiModalProcessor( |
| 443 | BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo] |
| 444 | ): |
| 445 | def get_dummy_inputs( |
| 446 | self, |
| 447 | seq_len: int, |
| 448 | mm_counts: Mapping[str, int], |
| 449 | mm_options: MultiModalDummyOptions, |
| 450 | ) -> ProcessorInputs: |
| 451 | builder = self.dummy_inputs |
| 452 | tokenizer = self.info.get_tokenizer() |
| 453 | |
| 454 | dummy_mm_data = builder.get_dummy_mm_data(seq_len, mm_counts, mm_options) |
| 455 | dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) |
| 456 | |
| 457 | num_audios = mm_counts.get("audio", 0) |
| 458 | if num_audios > 0: |
| 459 | prompt = cached_encode( |
| 460 | tokenizer, |
| 461 | AUDIO_PLACEHOLDER * num_audios, |
| 462 | add_special_tokens=False, |
| 463 | ) |
| 464 | else: |
| 465 | prompt = cached_encode( |
| 466 | tokenizer, |
| 467 | "\n", |
| 468 | add_special_tokens=False, |
| 469 | ) |
| 470 | |
| 471 | return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items) |
| 472 | |
| 473 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 474 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 475 | |
| 476 | def _get_hf_mm_inputs( |
| 477 | self, |
| 478 | mm_items: MultiModalDataItems, |
| 479 | hf_kwargs: Mapping[str, object], |
| 480 | ) -> HFMultiModalInputs: |
| 481 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 482 | if hf_inputs.hf_data: |
| 483 | hf_inputs.hf_data["audio"] = _get_audios_from_mm_data(hf_inputs.hf_data) |
| 484 | |
| 485 | return hf_inputs |
| 486 | |
| 487 | def _postprocess_hf_mm_data( |
| 488 | self, |
| 489 | mm_data: Mapping[str, object], |
| 490 | hf_processor_mm_kwargs: Mapping[str, object], |
| 491 | processed_data: BatchFeature, |
| 492 | ) -> BatchFeature: |
| 493 | audios = _get_audios_from_mm_data(mm_data) |
| 494 | if audios: |
| 495 | processed_data = _add_vllm_audio_metadata(processed_data, len(audios)) |
| 496 | |
| 497 | return processed_data |
| 498 | |
| 499 | def _get_mm_fields_config( |
| 500 | self, |
| 501 | hf_inputs: BatchFeature, |
| 502 | hf_processor_mm_kwargs: Mapping[str, object], |
| 503 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 504 | return _mtd_field_config(hf_inputs) |
| 505 | |
| 506 | def _get_prompt_updates( |
| 507 | self, |
| 508 | mm_items: MultiModalDataItems, |
| 509 | hf_processor_mm_kwargs: Mapping[str, object], |
| 510 | out_mm_kwargs: MultiModalKwargsItems, |
| 511 | ) -> Sequence[PromptUpdate]: |
| 512 | out_mm_data = out_mm_kwargs.get_data() |
| 513 | audio_token_lengths_tensor = out_mm_data.get("audio_token_lengths") |
| 514 | if audio_token_lengths_tensor is None: |
| 515 | audio_embeds = out_mm_data.get("audio_embeds") |
| 516 | if audio_embeds is None: |
| 517 | audio_token_lengths: list[int] = [] |
| 518 | else: |
| 519 | audio_token_lengths = [ |
| 520 | int(audio_embed.shape[0]) |
| 521 | for audio_embed in _as_audio_embedding_list(audio_embeds) |
| 522 | ] |
| 523 | else: |
| 524 | if not isinstance(audio_token_lengths_tensor, torch.Tensor): |
| 525 | raise TypeError( |
| 526 | "`audio_token_lengths` must be a torch.Tensor, got " |
| 527 | f"{type(audio_token_lengths_tensor)!r}." |
| 528 | ) |
| 529 | audio_token_lengths = [ |
| 530 | int(length) for length in audio_token_lengths_tensor.tolist() |
| 531 | ] |
| 532 | processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 533 | tokenizer = self.info.get_tokenizer() |
| 534 | audio_start_id = _get_required_token_id(tokenizer, processor.audio_start_token) |
| 535 | audio_token_id = int(processor.audio_token_id) |
| 536 | audio_end_id = _get_required_token_id(tokenizer, processor.audio_end_token) |
| 537 | |
| 538 | def get_num_tokens(item_idx: int) -> int: |
| 539 | if item_idx >= len(audio_token_lengths): |
| 540 | raise ValueError( |
| 541 | "Cannot determine the number of audio tokens for audio item " |
| 542 | f"{item_idx}." |
| 543 | ) |
| 544 | num_tokens = audio_token_lengths[item_idx] |
| 545 | if num_tokens <= 0: |
| 546 | raise ValueError("Audio input is too short to produce any tokens.") |
| 547 | return num_tokens |
| 548 | |
| 549 | def get_replacement(item_idx: int) -> PromptUpdateDetails: |
| 550 | num_tokens = get_num_tokens(item_idx) |
| 551 | audio_tokens = processor._audio_span_ids(num_tokens) |
| 552 | return PromptUpdateDetails.select_token_id( |
| 553 | [audio_start_id] + audio_tokens + [audio_end_id], |
| 554 | embed_token_id=audio_token_id, |
| 555 | ) |
| 556 | |
| 557 | return [ |
| 558 | PromptReplacement( |
| 559 | modality="audio", |
| 560 | target=cached_encode( |
| 561 | tokenizer, AUDIO_PLACEHOLDER, add_special_tokens=False |
| 562 | ), |
| 563 | replacement=get_replacement, |
| 564 | ), |
| 565 | ] |
| 566 | |
| 567 | |
| 568 | @MULTIMODAL_REGISTRY.register_processor( |
| 569 | MossTranscribeDiarizeMultiModalProcessor, |
| 570 | info=MossTranscribeDiarizeProcessingInfo, |
| 571 | dummy_inputs=MossTranscribeDiarizeDummyInputsBuilder, |
| 572 | ) |
| 573 | class MossTranscribeDiarizeForConditionalGeneration( |
| 574 | nn.Module, |
| 575 | SupportsMultiModal, |
| 576 | SupportsPP, |
| 577 | SupportsTranscription, |
| 578 | ): |
| 579 | supports_transcription = True |
| 580 | supports_transcription_only = True |
| 581 | supports_segment_timestamp = False |
| 582 | supports_diarized_transcription = True |
| 583 | supported_languages = ISO639_1_SUPPORTED_LANGS |
| 584 | hf_to_vllm_mapper = WeightsMapper( |
| 585 | orig_to_new_prefix={ |
| 586 | "language_model.layers.": "language_model.model.layers.", |
| 587 | "language_model.embed_tokens.": "language_model.model.embed_tokens.", |
| 588 | "language_model.norm.": "language_model.model.norm.", |
| 589 | "model.language_model.model.": "language_model.model.", |
| 590 | "model.language_model.lm_head.": "language_model.lm_head.", |
| 591 | "model.language_model.": "language_model.model.", |
| 592 | "model.whisper_encoder.": "whisper_encoder.", |
| 593 | "model.vq_adaptor.": "vq_adaptor.", |
| 594 | "model.": None, |
| 595 | }, |
| 596 | ) |
| 597 | |
| 598 | @classmethod |
| 599 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 600 | return AUDIO_PLACEHOLDER if modality.startswith("audio") else None |
| 601 | |
| 602 | @classmethod |
| 603 | def get_speech_to_text_config( |
| 604 | cls, |
| 605 | model_config: ModelConfig, |
| 606 | task_type: str, |
| 607 | ) -> SpeechToTextConfig: |
| 608 | processor = cached_processor_from_config(model_config) |
| 609 | return SpeechToTextConfig( |
| 610 | max_audio_clip_s=None, |
| 611 | sample_rate=processor.feature_extractor.sampling_rate, |
| 612 | min_energy_split_window_size=None, |
| 613 | ) |
| 614 | |
| 615 | @classmethod |
| 616 | def get_num_audio_tokens( |
| 617 | cls, |
| 618 | audio_duration_s: float, |
| 619 | stt_config: SpeechToTextConfig, |
| 620 | model_config: ModelConfig, |
| 621 | ) -> int | None: |
| 622 | processor = cached_processor_from_config(model_config) |
| 623 | num_samples = math.ceil(audio_duration_s * stt_config.sample_rate) |
| 624 | return _compute_total_audio_tokens( |
| 625 | num_samples, |
| 626 | processor.feature_extractor, |
| 627 | processor.audio_merge_size, |
| 628 | ) |
| 629 | |
| 630 | @classmethod |
| 631 | def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: |
| 632 | stt_config = stt_params.stt_config |
| 633 | question = stt_params.request_prompt or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT |
| 634 | question = question.strip() or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT |
| 635 | prompt = ( |
| 636 | "<|im_start|>system\n" |
| 637 | "You are a helpful assistant.<|im_end|>\n" |
| 638 | f"<|im_start|>user\n{AUDIO_PLACEHOLDER}\n" |
| 639 | f"{question}<|im_end|>\n" |
| 640 | "<|im_start|>assistant\n" |
| 641 | ) |
| 642 | return TextPrompt( |
| 643 | prompt=prompt, |
| 644 | multi_modal_data={"audio": (stt_params.audio, stt_config.sample_rate)}, |
| 645 | ) |
| 646 | |
| 647 | @classmethod |
| 648 | def post_process_output(cls, text: str) -> str: |
| 649 | return text.strip() |
| 650 | |
| 651 | @classmethod |
| 652 | def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegment]: |
| 653 | """Parse MOSS's canonical ``[start][Sxx]text[end]`` transcript.""" |
| 654 | headers: list[tuple[re.Match[str], float, str]] = [] |
| 655 | for match in _MOSS_DIARIZED_HEADER_RE.finditer(text): |
| 656 | start = parse_diarized_timestamp(match["start"]) |
| 657 | speaker = parse_diarized_speaker(match["speaker"]) |
| 658 | if start is not None and speaker is not None: |
| 659 | headers.append((match, start, speaker)) |
| 660 | |
| 661 | if not headers: |
| 662 | return [] |
| 663 | |
| 664 | segments: list[DiarizedTranscriptionSegment] = [] |
| 665 | for index, (header, start, speaker) in enumerate(headers): |
| 666 | next_header_start = ( |
| 667 | headers[index + 1][0].start() if index + 1 < len(headers) else len(text) |
| 668 | ) |
| 669 | body = text[header.end() : next_header_start] |
| 670 | end_match = _MOSS_DIARIZED_END_RE.search(body) |
| 671 | if end_match is None: |
| 672 | return [] |
| 673 | |
| 674 | end = parse_diarized_timestamp(end_match["end"]) |
| 675 | if end is None or end < start: |
| 676 | return [] |
| 677 | |
| 678 | segment_text = body[: end_match.start()].strip() |
| 679 | if segment_text: |
| 680 | segments.append( |
| 681 | DiarizedTranscriptionSegment( |
| 682 | start=start, |
| 683 | end=end, |
| 684 | speaker=speaker, |
| 685 | text=segment_text, |
| 686 | ) |
| 687 | ) |
| 688 | |
| 689 | return segments |
| 690 | |
| 691 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 692 | super().__init__() |
| 693 | self.config = vllm_config.model_config.hf_config |
| 694 | self.dtype = vllm_config.model_config.dtype |
| 695 | |
| 696 | with self._mark_tower_model(vllm_config, "audio"): |
| 697 | self.whisper_encoder = MossTranscribeDiarizeWhisperEncoder( |
| 698 | vllm_config=vllm_config, |
| 699 | prefix=maybe_prefix(prefix, "whisper_encoder"), |
| 700 | ) |
| 701 | self.vq_adaptor = MossTranscribeDiarizeVQAdaptor( |
| 702 | input_dim=int(self.config.adaptor_input_dim), |
| 703 | hidden_size=int(self.config.text_config.hidden_size), |
| 704 | eps=float(self.config.text_config.rms_norm_eps), |
| 705 | ) |
| 706 | |
| 707 | with self._mark_language_model(vllm_config): |
| 708 | self.language_model = init_vllm_registered_model( |
| 709 | vllm_config=vllm_config, |
| 710 | hf_config=self.config.text_config, |
| 711 | prefix=maybe_prefix(prefix, "language_model"), |
| 712 | architectures=["Qwen3ForCausalLM"], |
| 713 | ) |
| 714 | |
| 715 | self.make_empty_intermediate_tensors = ( |
| 716 | self.language_model.make_empty_intermediate_tensors |
| 717 | ) |
| 718 | |
| 719 | def _time_merge(self, features: torch.Tensor) -> torch.Tensor: |
| 720 | batch, seq_len, dim = features.shape |
| 721 | merge_size = int(self.config.audio_merge_size) |
| 722 | seq_len_trim = (seq_len // merge_size) * merge_size |
| 723 | return features[:, :seq_len_trim, :].reshape( |
| 724 | batch, |
| 725 | seq_len_trim // merge_size, |
| 726 | dim * merge_size, |
| 727 | ) |
| 728 | |
| 729 | def _parse_and_validate_audio_input( |
| 730 | self, |
| 731 | **kwargs: object, |
| 732 | ) -> MossTranscribeDiarizeInputs | None: |
| 733 | input_features = kwargs.pop("input_features", None) |
| 734 | audio_embeds = kwargs.pop("audio_embeds", None) |
| 735 | audio_feature_lengths = kwargs.pop("audio_feature_lengths", None) |
| 736 | audio_chunk_counts = kwargs.pop("audio_chunk_counts", None) |
| 737 | if input_features is None and audio_embeds is None: |
| 738 | return None |
| 739 | if audio_embeds is not None: |
| 740 | return MossTranscribeDiarizeEmbeddingInputs( |
| 741 | type="audio_embeds", |
| 742 | audio_embeds=_as_audio_embedding_list(audio_embeds), |
| 743 | ) |
| 744 | return MossTranscribeDiarizeAudioInputs( |
| 745 | type="audio_features", |
| 746 | input_features=input_features, |
| 747 | audio_feature_lengths=audio_feature_lengths, |
| 748 | audio_chunk_counts=audio_chunk_counts, |
| 749 | ) |
| 750 | |
| 751 | def _process_audio_input( |
| 752 | self, audio_input: MossTranscribeDiarizeInputs |
| 753 | ) -> list[torch.Tensor]: |
| 754 | if audio_input["type"] == "audio_embeds": |
| 755 | return list(audio_input["audio_embeds"]) |
| 756 | |
| 757 | input_features = audio_input["input_features"] |
| 758 | audio_feature_lengths = audio_input["audio_feature_lengths"] |
| 759 | audio_chunk_counts = audio_input["audio_chunk_counts"] |
| 760 | if input_features is None or audio_feature_lengths is None: |
| 761 | raise ValueError( |
| 762 | "MOSS-Transcribe-Diarize audio inputs require both " |
| 763 | "`input_features` and `audio_feature_lengths`." |
| 764 | ) |
| 765 | if audio_feature_lengths.numel() != input_features.shape[0]: |
| 766 | raise ValueError( |
| 767 | "`audio_feature_lengths` must contain one length per " |
| 768 | "`input_features` chunk: got " |
| 769 | f"{audio_feature_lengths.numel()} lengths for " |
| 770 | f"{input_features.shape[0]} chunks." |
| 771 | ) |
| 772 | if audio_chunk_counts is None: |
| 773 | audio_chunk_counts = audio_feature_lengths.new_tensor( |
| 774 | [input_features.shape[0]], |
| 775 | dtype=torch.long, |
| 776 | ) |
| 777 | else: |
| 778 | audio_chunk_counts = audio_chunk_counts.to(dtype=torch.long) |
| 779 | if audio_chunk_counts.numel() == 0: |
| 780 | raise ValueError("`audio_chunk_counts` must contain at least one item.") |
| 781 | if torch.any(audio_chunk_counts <= 0): |
| 782 | raise ValueError("`audio_chunk_counts` must contain positive counts.") |
| 783 | num_audio_chunks = int(audio_chunk_counts.sum().item()) |
| 784 | if num_audio_chunks != input_features.shape[0]: |
| 785 | raise ValueError( |
| 786 | "`audio_chunk_counts` must sum to the number of input feature chunks: " |
| 787 | f"got {num_audio_chunks} chunks for " |
| 788 | f"{input_features.shape[0]} input chunks." |
| 789 | ) |
| 790 | |
| 791 | features = self.whisper_encoder(input_features, audio_feature_lengths) |
| 792 | merged = self._time_merge(features.to(dtype=self.dtype)) |
| 793 | projected = self.vq_adaptor(merged).squeeze(0) |
| 794 | |
| 795 | audio_chunk_offsets = torch.cumsum(audio_chunk_counts, dim=0) |
| 796 | audio_chunk_offsets = torch.cat( |
| 797 | [audio_chunk_offsets.new_zeros(1), audio_chunk_offsets] |
| 798 | ) |
| 799 | tokens_per_item = [ |
| 800 | int(audio_feature_lengths[start:end].sum().item()) |
| 801 | for start, end in zip( |
| 802 | audio_chunk_offsets[:-1].tolist(), |
| 803 | audio_chunk_offsets[1:].tolist(), |
| 804 | ) |
| 805 | ] |
| 806 | if any(num_tokens <= 0 for num_tokens in tokens_per_item): |
| 807 | raise ValueError("Audio input is too short to produce any tokens.") |
| 808 | return list(projected.split(tokens_per_item, dim=0)) |
| 809 | |
| 810 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 811 | audio_input = self._parse_and_validate_audio_input(**kwargs) |
| 812 | if audio_input is None: |
| 813 | return [] |
| 814 | return self._process_audio_input(audio_input) |
| 815 | |
| 816 | def embed_input_ids( |
| 817 | self, |
| 818 | input_ids: torch.Tensor, |
| 819 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 820 | *, |
| 821 | is_multimodal: torch.Tensor | None = None, |
| 822 | ) -> torch.Tensor: |
| 823 | inputs_embeds = self.language_model.embed_input_ids(input_ids) |
| 824 | if not multimodal_embeddings: |
| 825 | return inputs_embeds |
| 826 | return _merge_multimodal_embeddings( |
| 827 | inputs_embeds=inputs_embeds, |
| 828 | multimodal_embeddings=multimodal_embeddings, |
| 829 | is_multimodal=_require_is_multimodal(is_multimodal), |
| 830 | ) |
| 831 | |
| 832 | def forward( |
| 833 | self, |
| 834 | input_ids: torch.Tensor | None, |
| 835 | positions: torch.Tensor, |
| 836 | intermediate_tensors: IntermediateTensors | None = None, |
| 837 | inputs_embeds: torch.Tensor | None = None, |
| 838 | **kwargs: object, |
| 839 | ) -> torch.Tensor | IntermediateTensors: |
| 840 | if intermediate_tensors is not None: |
| 841 | inputs_embeds = None |
| 842 | return self.language_model( |
| 843 | input_ids=input_ids, |
| 844 | positions=positions, |
| 845 | intermediate_tensors=intermediate_tensors, |
| 846 | inputs_embeds=inputs_embeds, |
| 847 | ) |
| 848 | |
| 849 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 850 | return self.language_model.compute_logits(hidden_states) |
| 851 | |
| 852 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 853 | loader = AutoWeightsLoader(self) |
| 854 | return loader.load_weights( |
| 855 | weights, |
| 856 | mapper=self.hf_to_vllm_mapper, |
| 857 | ) |