[Model] Support MOSS-Transcribe-Diarize (#47729)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Canlin Guo committed
Jul 8, 2026 at 19:05 UTC
285c08c036880d404e35b175008934a6aeeb4832
7 files changed
+925
docs/models/supported_models.md
+2
@@ -569,6 +569,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
569
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
570
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
571
| `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
572
+| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
573
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
574
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
575
| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
@@ -671,6 +672,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
672
| `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
673
| `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ |
674
| `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
675
+| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
676
| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ |
677
| `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ |
678
| `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+5
@@ -1127,6 +1127,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1127
},
1128
trust_remote_code=True,
1129
),
1130
+ "MossTranscribeDiarizeForConditionalGeneration": _HfExamplesInfo(
1131
+ "OpenMOSS-Team/MOSS-Transcribe-Diarize",
1132
+ trust_remote_code=True,
1133
+ is_available_online=False,
1134
+ ),
1135
"HfMoondream": _HfExamplesInfo(
1136
"moondream/moondream3-preview",
1137
tokenizer="moondream/starmie-v1",
vllm/model_executor/models/moss_transcribe_diarize.py
new
+801
@@ -0,0 +1,801 @@
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 torch
18
+from torch import nn
19
+from transformers import BatchFeature
20
+
21
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
22
+from vllm.config.multimodal import BaseDummyOptions
23
+from vllm.config.speech_to_text import SpeechToTextParams
24
+from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt
25
+from vllm.model_executor.models.interfaces import (
26
+ MultiModalEmbeddings,
27
+ SupportsMultiModal,
28
+ SupportsPP,
29
+ SupportsTranscription,
30
+ _require_is_multimodal,
31
+)
32
+from vllm.model_executor.models.utils import (
33
+ AutoWeightsLoader,
34
+ WeightsMapper,
35
+ _merge_multimodal_embeddings,
36
+ init_vllm_registered_model,
37
+ maybe_prefix,
38
+)
39
+from vllm.model_executor.models.whisper import (
40
+ WhisperEncoder,
41
+ _create_fake_bias_for_k_proj,
42
+)
43
+from vllm.model_executor.models.whisper_utils import ISO639_1_SUPPORTED_LANGS
44
+from vllm.multimodal import MULTIMODAL_REGISTRY
45
+from vllm.multimodal.inputs import (
46
+ AudioItem,
47
+ MultiModalFieldConfig,
48
+ MultiModalKwargsItems,
49
+)
50
+from vllm.multimodal.parse import (
51
+ DictEmbeddingItems,
52
+ ModalityDataItems,
53
+ MultiModalDataItems,
54
+ MultiModalDataParser,
55
+)
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 ProcessorInputs
65
+from vllm.sequence import IntermediateTensors
66
+from vllm.transformers_utils.processor import cached_processor_from_config
67
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
68
+
69
+WHISPER_ENCODER_STRIDE = 2
70
+
71
+AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>"
72
+
73
+DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT = (
74
+ "请将音频转写为文本，每一段需以起始时间戳和说话人编号"
75
+ "（[S01]、[S02]、[S03]…）开头，正文为对应的语音内容，"
76
+ "并在段末标注结束时间戳，以清晰标明该段语音范围。"
77
+)
78
+
79
+
80
+class MossTranscribeDiarizeAudioInputs(TensorSchema):
81
+ """
82
+ Dimensions:
83
+ - c: Audio chunks
84
+ - m: Mel bins
85
+ - f: Mel frames
86
+ - n: Number of audio items
87
+ """
88
+
89
+ type: Literal["audio_features"] = "audio_features"
90
+
91
+ input_features: Annotated[
92
+ torch.Tensor | None,
93
+ TensorShape("c", "m", "f"),
94
+ ]
95
+ audio_feature_lengths: Annotated[
96
+ torch.Tensor | None,
97
+ TensorShape("c"),
98
+ ]
99
+ audio_chunk_counts: Annotated[
100
+ torch.Tensor | None,
101
+ TensorShape("n"),
102
+ ]
103
+
104
+
105
+class MossTranscribeDiarizeEmbeddingInputs(TensorSchema):
106
+ """
107
+ Dimensions:
108
+ - n: Number of audio items
109
+ - t: Number of audio tokens
110
+ - h: Hidden size
111
+ """
112
+
113
+ type: Literal["audio_embeds"] = "audio_embeds"
114
+
115
+ audio_embeds: Annotated[
116
+ list[torch.Tensor],
117
+ TensorShape("n", "t", "h", dynamic_dims={"t"}),
118
+ ]
119
+
120
+
121
+MossTranscribeDiarizeInputs: TypeAlias = (
122
+ MossTranscribeDiarizeAudioInputs | MossTranscribeDiarizeEmbeddingInputs
123
+)
124
+
125
+
126
+def _compute_total_audio_tokens(
127
+ num_samples: int,
128
+ feature_extractor: Any,
129
+ audio_merge_size: int,
130
+) -> int:
131
+ if num_samples <= 0:
132
+ return 0
133
+
134
+ n_samples = int(feature_extractor.n_samples)
135
+ stride = (
136
+ int(feature_extractor.hop_length)
137
+ * WHISPER_ENCODER_STRIDE
138
+ * int(audio_merge_size)
139
+ )
140
+ total = 0
141
+ for start in range(0, num_samples, n_samples):
142
+ chunk_samples = min(n_samples, num_samples - start)
143
+ total += (chunk_samples - 1) // stride + 1
144
+ return total
145
+
146
+
147
+def _get_max_audio_samples(feature_extractor: Any) -> int:
148
+ if hasattr(feature_extractor, "chunk_length"):
149
+ return int(feature_extractor.chunk_length * feature_extractor.sampling_rate)
150
+ return int(feature_extractor.n_samples)
151
+
152
+
153
+def _as_audio_embedding_list(audio_embeds: object) -> list[torch.Tensor]:
154
+ if isinstance(audio_embeds, torch.Tensor):
155
+ if audio_embeds.ndim == 2:
156
+ return [audio_embeds]
157
+ if audio_embeds.ndim == 3:
158
+ return list(audio_embeds.unbind(dim=0))
159
+ raise ValueError(
160
+ f"`audio_embeds` must be a 2D or 3D tensor, got {audio_embeds.ndim}D."
161
+ )
162
+
163
+ if isinstance(audio_embeds, (list, tuple)) and all(
164
+ isinstance(audio_embed, torch.Tensor) for audio_embed in audio_embeds
165
+ ):
166
+ return list(audio_embeds)
167
+
168
+ raise TypeError(
169
+ "`audio_embeds` must be a torch.Tensor or a list of torch.Tensor "
170
+ f"objects, got {type(audio_embeds)!r}."
171
+ )
172
+
173
+
174
+def _get_required_token_id(tokenizer: Any, token: str) -> int:
175
+ token_id = tokenizer.convert_tokens_to_ids(token)
176
+ if token_id is None or (isinstance(token_id, int) and token_id < 0):
177
+ raise ValueError(f"Tokenizer is missing required token {token!r}.")
178
+ return int(token_id)
179
+
180
+
181
+def _get_audios_from_mm_data(mm_data: Mapping[str, object]) -> list[Any]:
182
+ mm_data_dict = dict(mm_data)
183
+ audio_data = mm_data_dict.pop("audios", mm_data_dict.pop("audio", []))
184
+ if isinstance(audio_data, list):
185
+ audios = audio_data
186
+ elif isinstance(audio_data, tuple) and not (
187
+ len(audio_data) == 2 and isinstance(audio_data[1], (int, float))
188
+ ):
189
+ audios = list(audio_data)
190
+ elif audio_data is None:
191
+ audios = []
192
+ else:
193
+ audios = [audio_data]
194
+
195
+ audio_arrays: list[Any] = []
196
+ for audio in audios:
197
+ if isinstance(audio, (tuple, list)) and len(audio) == 2:
198
+ audio = audio[0]
199
+ audio_arrays.append(audio)
200
+ return audio_arrays
201
+
202
+
203
+def _add_vllm_audio_metadata(
204
+ processed: BatchFeature,
205
+ num_audios: int,
206
+) -> BatchFeature:
207
+ audio_feature_lengths = processed["audio_feature_lengths"]
208
+ if not isinstance(audio_feature_lengths, torch.Tensor):
209
+ audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long)
210
+ audio_feature_lengths = audio_feature_lengths.to(dtype=torch.long)
211
+
212
+ audio_chunk_mapping = processed.get("audio_chunk_mapping")
213
+ if audio_chunk_mapping is None:
214
+ if num_audios == 1:
215
+ audio_chunk_mapping = torch.zeros_like(audio_feature_lengths)
216
+ elif audio_feature_lengths.numel() == num_audios:
217
+ audio_chunk_mapping = torch.arange(num_audios, dtype=torch.long)
218
+ else:
219
+ raise ValueError(
220
+ "The MOSS processor did not return `audio_chunk_mapping`, and "
221
+ "the chunk-to-audio mapping cannot be inferred."
222
+ )
223
+ elif not isinstance(audio_chunk_mapping, torch.Tensor):
224
+ audio_chunk_mapping = torch.tensor(audio_chunk_mapping, dtype=torch.long)
225
+ audio_chunk_mapping = audio_chunk_mapping.to(dtype=torch.long)
226
+ if audio_chunk_mapping.numel() != audio_feature_lengths.numel():
227
+ raise ValueError(
228
+ "`audio_chunk_mapping` must contain one item per audio chunk: got "
229
+ f"{audio_chunk_mapping.numel()} mappings for "
230
+ f"{audio_feature_lengths.numel()} chunks."
231
+ )
232
+ if audio_chunk_mapping.numel() > 0 and (
233
+ audio_chunk_mapping.min().item() < 0
234
+ or audio_chunk_mapping.max().item() >= num_audios
235
+ ):
236
+ raise ValueError("`audio_chunk_mapping` contains an out-of-range audio index.")
237
+
238
+ audio_chunk_counts = torch.bincount(
239
+ audio_chunk_mapping.cpu(),
240
+ minlength=num_audios,
241
+ ).to(dtype=torch.long)
242
+ audio_token_lengths = torch.zeros(
243
+ num_audios,
244
+ dtype=torch.long,
245
+ device=audio_feature_lengths.device,
246
+ )
247
+ audio_token_lengths.scatter_add_(
248
+ 0,
249
+ audio_chunk_mapping.to(device=audio_feature_lengths.device),
250
+ audio_feature_lengths,
251
+ )
252
+
253
+ processed["audio_feature_lengths"] = audio_feature_lengths
254
+ processed["audio_chunk_counts"] = audio_chunk_counts
255
+ processed["audio_token_lengths"] = audio_token_lengths
256
+ return processed
257
+
258
+
259
+class MossTranscribeDiarizeWhisperEncoder(WhisperEncoder):
260
+ packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
261
+ hf_to_vllm_mapper = WeightsMapper(
262
+ orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."},
263
+ orig_to_new_stacked={
264
+ ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
265
+ ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
266
+ ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
267
+ },
268
+ )
269
+
270
+ def __init__(
271
+ self,
272
+ *,
273
+ vllm_config: VllmConfig,
274
+ prefix: str = "",
275
+ ) -> None:
276
+ super().__init__(
277
+ vllm_config=vllm_config.with_hf_config(
278
+ vllm_config.model_config.hf_config.audio_config
279
+ ),
280
+ prefix=prefix,
281
+ )
282
+ self.audio_merge_size = int(vllm_config.model_config.hf_config.audio_merge_size)
283
+ self.max_encoder_batch: int | None = None
284
+
285
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
286
+ weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
287
+ loader = AutoWeightsLoader(self)
288
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
289
+
290
+ def forward(
291
+ self,
292
+ input_features: torch.Tensor,
293
+ audio_feature_lengths: torch.Tensor,
294
+ ) -> torch.Tensor:
295
+ if input_features.numel() == 0:
296
+ return input_features.new_empty((1, 0, self.conv1.out_channels))
297
+ device = self.conv1.weight.device
298
+ dtype = self.conv1.weight.dtype
299
+ input_features = input_features.to(device=device, dtype=dtype)
300
+ audio_feature_lengths = audio_feature_lengths.to(device=device)
301
+ batch_size = self.max_encoder_batch or input_features.shape[0]
302
+ encoded_parts: list[torch.Tensor] = []
303
+
304
+ for start in range(0, input_features.shape[0], batch_size):
305
+ encoded_parts.append(
306
+ super().forward([input_features[start : start + batch_size]])
307
+ )
308
+
309
+ hidden = torch.cat(encoded_parts, dim=0)
310
+ chunks = [
311
+ hidden[idx : idx + 1, : int(token_len.item()) * self.audio_merge_size]
312
+ for idx, token_len in enumerate(audio_feature_lengths)
313
+ ]
314
+ return torch.cat(chunks, dim=1)
315
+
316
+
317
+class MossTranscribeDiarizeVQAdaptor(nn.Module):
318
+ def __init__(self, input_dim: int, hidden_size: int, eps: float) -> None:
319
+ super().__init__()
320
+ self.layers = nn.Sequential(
321
+ nn.Linear(input_dim, hidden_size, bias=True),
322
+ nn.SiLU(),
323
+ nn.Linear(hidden_size, hidden_size, bias=True),
324
+ nn.LayerNorm(hidden_size, eps=eps, bias=True),
325
+ )
326
+
327
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
328
+ return self.layers(x)
329
+
330
+
331
+def _mtd_field_config(
332
+ hf_inputs: Mapping[str, torch.Tensor],
333
+) -> Mapping[str, MultiModalFieldConfig]:
334
+ fields: dict[str, MultiModalFieldConfig] = {}
335
+ if "audio_embeds" in hf_inputs:
336
+ fields["audio_embeds"] = MultiModalFieldConfig.batched("audio")
337
+ if "audio_chunk_counts" in hf_inputs:
338
+ audio_chunk_counts = hf_inputs["audio_chunk_counts"]
339
+ fields.update(
340
+ input_features=MultiModalFieldConfig.flat_from_sizes(
341
+ "audio",
342
+ audio_chunk_counts,
343
+ ),
344
+ audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes(
345
+ "audio",
346
+ audio_chunk_counts,
347
+ ),
348
+ audio_chunk_counts=MultiModalFieldConfig.batched("audio"),
349
+ audio_token_lengths=MultiModalFieldConfig.batched("audio"),
350
+ )
351
+ return fields
352
+
353
+
354
+class MossTranscribeDiarizeMultiModalDataParser(MultiModalDataParser):
355
+ def _parse_audio_data(
356
+ self,
357
+ data: dict[str, torch.Tensor] | ModalityData[AudioItem],
358
+ ) -> ModalityDataItems[Any, Any] | None:
359
+ if isinstance(data, dict):
360
+ return DictEmbeddingItems(
361
+ data,
362
+ modality="audio",
363
+ required_fields={"audio_embeds"},
364
+ fields_factory=_mtd_field_config,
365
+ )
366
+
367
+ return super()._parse_audio_data(data)
368
+
369
+
370
+class MossTranscribeDiarizeProcessingInfo(BaseProcessingInfo):
371
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
372
+ return {"audio": 1}
373
+
374
+ def get_hf_processor(self, **kwargs: object) -> Any:
375
+ return self.ctx.get_hf_processor(**kwargs)
376
+
377
+ def get_feature_extractor(self, **kwargs: object) -> Any:
378
+ return self.get_hf_processor(**kwargs).feature_extractor
379
+
380
+ def get_data_parser(self) -> MultiModalDataParser:
381
+ feature_extractor = self.get_feature_extractor()
382
+ return MossTranscribeDiarizeMultiModalDataParser(
383
+ target_sr=feature_extractor.sampling_rate,
384
+ target_channels=1,
385
+ expected_hidden_size=self._get_expected_hidden_size(),
386
+ )
387
+
388
+ def get_mm_max_tokens_per_item(
389
+ self,
390
+ seq_len: int,
391
+ mm_counts: Mapping[str, int],
392
+ ) -> Mapping[str, int]:
393
+ if mm_counts.get("audio", 0) <= 0:
394
+ return {}
395
+
396
+ feature_extractor = self.get_feature_extractor()
397
+ max_audio_samples = _get_max_audio_samples(feature_extractor)
398
+ max_audio_tokens = _compute_total_audio_tokens(
399
+ max_audio_samples,
400
+ feature_extractor,
401
+ self.get_hf_processor().audio_merge_size,
402
+ )
403
+ return {"audio": min(seq_len, max_audio_tokens)}
404
+
405
+
406
+class MossTranscribeDiarizeDummyInputsBuilder(
407
+ BaseDummyInputsBuilder[MossTranscribeDiarizeProcessingInfo]
408
+):
409
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
410
+ return AUDIO_PLACEHOLDER * mm_counts.get("audio", 0)
411
+
412
+ def get_dummy_mm_data(
413
+ self,
414
+ seq_len: int,
415
+ mm_counts: Mapping[str, int],
416
+ mm_options: Mapping[str, BaseDummyOptions],
417
+ ) -> MultiModalDataDict:
418
+ num_audios = mm_counts.get("audio", 0)
419
+ if num_audios == 0:
420
+ return {}
421
+
422
+ feature_extractor = self.info.get_feature_extractor()
423
+ return {
424
+ "audio": self._get_dummy_audios(
425
+ length=_get_max_audio_samples(feature_extractor),
426
+ num_audios=num_audios,
427
+ overrides=mm_options.get("audio"),
428
+ )
429
+ }
430
+
431
+ def get_dummy_processor_inputs(
432
+ self,
433
+ seq_len: int,
434
+ mm_counts: Mapping[str, int],
435
+ mm_options: Mapping[str, BaseDummyOptions],
436
+ ) -> ProcessorInputs:
437
+ dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
438
+ dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
439
+ num_audios = mm_counts.get("audio", 0)
440
+ tokenizer = self.info.get_tokenizer()
441
+ prompt = tokenizer.encode(
442
+ AUDIO_PLACEHOLDER * num_audios,
443
+ add_special_tokens=False,
444
+ ) or tokenizer.encode(
445
+ "\n",
446
+ add_special_tokens=False,
447
+ )
448
+ return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items)
449
+
450
+
451
+class MossTranscribeDiarizeMultiModalProcessor(
452
+ BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo]
453
+):
454
+ def _call_hf_processor(
455
+ self,
456
+ prompt: str,
457
+ mm_data: Mapping[str, object],
458
+ mm_kwargs: Mapping[str, object],
459
+ tok_kwargs: Mapping[str, object],
460
+ ) -> BatchFeature:
461
+ tokenizer = self.info.get_tokenizer()
462
+ audios = _get_audios_from_mm_data(mm_data)
463
+ if not audios:
464
+ input_ids = tokenizer.encode(
465
+ prompt,
466
+ add_special_tokens=tok_kwargs.get("add_special_tokens", False),
467
+ )
468
+ return BatchFeature({"input_ids": [input_ids]}, tensor_type="pt")
469
+
470
+ processed = self.info.ctx.call_hf_processor(
471
+ self.info.get_hf_processor(**mm_kwargs),
472
+ dict(text=prompt, audio=audios),
473
+ dict(**mm_kwargs, **tok_kwargs),
474
+ )
475
+ return _add_vllm_audio_metadata(processed, len(audios))
476
+
477
+ def _hf_processor_applies_updates(
478
+ self,
479
+ prompt_text: str,
480
+ mm_items: MultiModalDataItems,
481
+ hf_processor_mm_kwargs: Mapping[str, object],
482
+ tokenization_kwargs: Mapping[str, object],
483
+ ) -> bool:
484
+ return mm_items.get_count("audio", strict=False) > 0
485
+
486
+ def _get_mm_fields_config(
487
+ self,
488
+ hf_inputs: BatchFeature,
489
+ hf_processor_mm_kwargs: Mapping[str, object],
490
+ ) -> Mapping[str, MultiModalFieldConfig]:
491
+ return _mtd_field_config(hf_inputs)
492
+
493
+ def _get_prompt_updates(
494
+ self,
495
+ mm_items: MultiModalDataItems,
496
+ hf_processor_mm_kwargs: Mapping[str, object],
497
+ out_mm_kwargs: MultiModalKwargsItems,
498
+ ) -> Sequence[PromptUpdate]:
499
+ out_mm_data = out_mm_kwargs.get_data()
500
+ audio_token_lengths_tensor = out_mm_data.get("audio_token_lengths")
501
+ if audio_token_lengths_tensor is None:
502
+ audio_embeds = out_mm_data.get("audio_embeds")
503
+ if audio_embeds is None:
504
+ audio_token_lengths: list[int] = []
505
+ else:
506
+ audio_token_lengths = [
507
+ int(audio_embed.shape[0])
508
+ for audio_embed in _as_audio_embedding_list(audio_embeds)
509
+ ]
510
+ else:
511
+ if not isinstance(audio_token_lengths_tensor, torch.Tensor):
512
+ raise TypeError(
513
+ "`audio_token_lengths` must be a torch.Tensor, got "
514
+ f"{type(audio_token_lengths_tensor)!r}."
515
+ )
516
+ audio_token_lengths = [
517
+ int(length) for length in audio_token_lengths_tensor.tolist()
518
+ ]
519
+ processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
520
+ tokenizer = self.info.get_tokenizer()
521
+ audio_start_id = _get_required_token_id(tokenizer, processor.audio_start_token)
522
+ audio_token_id = int(processor.audio_token_id)
523
+ audio_end_id = _get_required_token_id(tokenizer, processor.audio_end_token)
524
+
525
+ def get_num_tokens(item_idx: int) -> int:
526
+ if item_idx >= len(audio_token_lengths):
527
+ raise ValueError(
528
+ "Cannot determine the number of audio tokens for audio item "
529
+ f"{item_idx}."
530
+ )
531
+ num_tokens = audio_token_lengths[item_idx]
532
+ if num_tokens <= 0:
533
+ raise ValueError("Audio input is too short to produce any tokens.")
534
+ return num_tokens
535
+
536
+ def get_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
537
+ num_tokens = get_num_tokens(item_idx)
538
+ audio_tokens = processor._audio_span_ids(num_tokens)
539
+ return PromptUpdateDetails.select_token_id(
540
+ [audio_start_id] + audio_tokens + [audio_end_id],
541
+ embed_token_id=audio_token_id,
542
+ )
543
+
544
+ return [
545
+ PromptReplacement(
546
+ modality="audio",
547
+ target=AUDIO_PLACEHOLDER,
548
+ replacement=get_replacement,
549
+ ),
550
+ ]
551
+
552
+
553
+@MULTIMODAL_REGISTRY.register_processor(
554
+ MossTranscribeDiarizeMultiModalProcessor,
555
+ info=MossTranscribeDiarizeProcessingInfo,
556
+ dummy_inputs=MossTranscribeDiarizeDummyInputsBuilder,
557
+)
558
+class MossTranscribeDiarizeForConditionalGeneration(
559
+ nn.Module,
560
+ SupportsMultiModal,
561
+ SupportsPP,
562
+ SupportsTranscription,
563
+):
564
+ supports_transcription = True
565
+ supports_transcription_only = True
566
+ supports_segment_timestamp = False
567
+ supported_languages = ISO639_1_SUPPORTED_LANGS
568
+ hf_to_vllm_mapper = WeightsMapper(
569
+ orig_to_new_prefix={
570
+ "language_model.layers.": "language_model.model.layers.",
571
+ "language_model.embed_tokens.": "language_model.model.embed_tokens.",
572
+ "language_model.norm.": "language_model.model.norm.",
573
+ "model.language_model.model.": "language_model.model.",
574
+ "model.language_model.lm_head.": "language_model.lm_head.",
575
+ "model.language_model.": "language_model.model.",
576
+ "model.whisper_encoder.": "whisper_encoder.",
577
+ "model.vq_adaptor.": "vq_adaptor.",
578
+ "model.": None,
579
+ },
580
+ )
581
+
582
+ @classmethod
583
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
584
+ return AUDIO_PLACEHOLDER if modality.startswith("audio") else None
585
+
586
+ @classmethod
587
+ def get_speech_to_text_config(
588
+ cls,
589
+ model_config: ModelConfig,
590
+ task_type: str,
591
+ ) -> SpeechToTextConfig:
592
+ processor = cached_processor_from_config(model_config)
593
+ return SpeechToTextConfig(
594
+ max_audio_clip_s=None,
595
+ sample_rate=processor.feature_extractor.sampling_rate,
596
+ min_energy_split_window_size=None,
597
+ )
598
+
599
+ @classmethod
600
+ def get_num_audio_tokens(
601
+ cls,
602
+ audio_duration_s: float,
603
+ stt_config: SpeechToTextConfig,
604
+ model_config: ModelConfig,
605
+ ) -> int | None:
606
+ processor = cached_processor_from_config(model_config)
607
+ num_samples = math.ceil(audio_duration_s * stt_config.sample_rate)
608
+ return _compute_total_audio_tokens(
609
+ num_samples,
610
+ processor.feature_extractor,
611
+ processor.audio_merge_size,
612
+ )
613
+
614
+ @classmethod
615
+ def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
616
+ stt_config = stt_params.stt_config
617
+ question = stt_params.request_prompt or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT
618
+ question = question.strip() or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT
619
+ prompt = (
620
+ "<|im_start|>system\n"
621
+ "You are a helpful assistant.<|im_end|>\n"
622
+ f"<|im_start|>user\n{AUDIO_PLACEHOLDER}\n"
623
+ f"{question}<|im_end|>\n"
624
+ "<|im_start|>assistant\n"
625
+ )
626
+ return TextPrompt(
627
+ prompt=prompt,
628
+ multi_modal_data={"audio": (stt_params.audio, stt_config.sample_rate)},
629
+ )
630
+
631
+ @classmethod
632
+ def post_process_output(cls, text: str) -> str:
633
+ return text.strip()
634
+
635
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
636
+ super().__init__()
637
+ self.config = vllm_config.model_config.hf_config
638
+ self.dtype = vllm_config.model_config.dtype
639
+
640
+ with self._mark_tower_model(vllm_config, "audio"):
641
+ self.whisper_encoder = MossTranscribeDiarizeWhisperEncoder(
642
+ vllm_config=vllm_config,
643
+ prefix=maybe_prefix(prefix, "whisper_encoder"),
644
+ )
645
+ self.vq_adaptor = MossTranscribeDiarizeVQAdaptor(
646
+ input_dim=int(self.config.adaptor_input_dim),
647
+ hidden_size=int(self.config.text_config.hidden_size),
648
+ eps=float(self.config.text_config.rms_norm_eps),
649
+ )
650
+
651
+ with self._mark_language_model(vllm_config):
652
+ self.language_model = init_vllm_registered_model(
653
+ vllm_config=vllm_config,
654
+ hf_config=self.config.text_config,
655
+ prefix=maybe_prefix(prefix, "language_model"),
656
+ architectures=["Qwen3ForCausalLM"],
657
+ )
658
+
659
+ self.make_empty_intermediate_tensors = (
660
+ self.language_model.make_empty_intermediate_tensors
661
+ )
662
+
663
+ def _time_merge(self, features: torch.Tensor) -> torch.Tensor:
664
+ batch, seq_len, dim = features.shape
665
+ merge_size = int(self.config.audio_merge_size)
666
+ seq_len_trim = (seq_len // merge_size) * merge_size
667
+ return features[:, :seq_len_trim, :].reshape(
668
+ batch,
669
+ seq_len_trim // merge_size,
670
+ dim * merge_size,
671
+ )
672
+
673
+ def _parse_and_validate_audio_input(
674
+ self,
675
+ **kwargs: object,
676
+ ) -> MossTranscribeDiarizeInputs | None:
677
+ input_features = kwargs.pop("input_features", None)
678
+ audio_embeds = kwargs.pop("audio_embeds", None)
679
+ audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
680
+ audio_chunk_counts = kwargs.pop("audio_chunk_counts", None)
681
+ if input_features is None and audio_embeds is None:
682
+ return None
683
+ if audio_embeds is not None:
684
+ return MossTranscribeDiarizeEmbeddingInputs(
685
+ type="audio_embeds",
686
+ audio_embeds=_as_audio_embedding_list(audio_embeds),
687
+ )
688
+ return MossTranscribeDiarizeAudioInputs(
689
+ type="audio_features",
690
+ input_features=input_features,
691
+ audio_feature_lengths=audio_feature_lengths,
692
+ audio_chunk_counts=audio_chunk_counts,
693
+ )
694
+
695
+ def _process_audio_input(
696
+ self, audio_input: MossTranscribeDiarizeInputs
697
+ ) -> list[torch.Tensor]:
698
+ if audio_input["type"] == "audio_embeds":
699
+ return list(audio_input["audio_embeds"])
700
+
701
+ input_features = audio_input["input_features"]
702
+ audio_feature_lengths = audio_input["audio_feature_lengths"]
703
+ audio_chunk_counts = audio_input["audio_chunk_counts"]
704
+ if input_features is None or audio_feature_lengths is None:
705
+ raise ValueError(
706
+ "MOSS-Transcribe-Diarize audio inputs require both "
707
+ "`input_features` and `audio_feature_lengths`."
708
+ )
709
+ if audio_feature_lengths.numel() != input_features.shape[0]:
710
+ raise ValueError(
711
+ "`audio_feature_lengths` must contain one length per "
712
+ "`input_features` chunk: got "
713
+ f"{audio_feature_lengths.numel()} lengths for "
714
+ f"{input_features.shape[0]} chunks."
715
+ )
716
+ if audio_chunk_counts is None:
717
+ audio_chunk_counts = audio_feature_lengths.new_tensor(
718
+ [input_features.shape[0]],
719
+ dtype=torch.long,
720
+ )
721
+ else:
722
+ audio_chunk_counts = audio_chunk_counts.to(dtype=torch.long)
723
+ if audio_chunk_counts.numel() == 0:
724
+ raise ValueError("`audio_chunk_counts` must contain at least one item.")
725
+ if torch.any(audio_chunk_counts <= 0):
726
+ raise ValueError("`audio_chunk_counts` must contain positive counts.")
727
+ num_audio_chunks = int(audio_chunk_counts.sum().item())
728
+ if num_audio_chunks != input_features.shape[0]:
729
+ raise ValueError(
730
+ "`audio_chunk_counts` must sum to the number of input feature chunks: "
731
+ f"got {num_audio_chunks} chunks for "
732
+ f"{input_features.shape[0]} input chunks."
733
+ )
734
+
735
+ features = self.whisper_encoder(input_features, audio_feature_lengths)
736
+ merged = self._time_merge(features.to(dtype=self.dtype))
737
+ projected = self.vq_adaptor(merged).squeeze(0)
738
+
739
+ audio_chunk_offsets = torch.cumsum(audio_chunk_counts, dim=0)
740
+ audio_chunk_offsets = torch.cat(
741
+ [audio_chunk_offsets.new_zeros(1), audio_chunk_offsets]
742
+ )
743
+ tokens_per_item = [
744
+ int(audio_feature_lengths[start:end].sum().item())
745
+ for start, end in zip(
746
+ audio_chunk_offsets[:-1].tolist(),
747
+ audio_chunk_offsets[1:].tolist(),
748
+ )
749
+ ]
750
+ if any(num_tokens <= 0 for num_tokens in tokens_per_item):
751
+ raise ValueError("Audio input is too short to produce any tokens.")
752
+ return list(projected.split(tokens_per_item, dim=0))
753
+
754
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
755
+ audio_input = self._parse_and_validate_audio_input(**kwargs)
756
+ if audio_input is None:
757
+ return []
758
+ return self._process_audio_input(audio_input)
759
+
760
+ def embed_input_ids(
761
+ self,
762
+ input_ids: torch.Tensor,
763
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
764
+ *,
765
+ is_multimodal: torch.Tensor | None = None,
766
+ ) -> torch.Tensor:
767
+ inputs_embeds = self.language_model.embed_input_ids(input_ids)
768
+ if not multimodal_embeddings:
769
+ return inputs_embeds
770
+ return _merge_multimodal_embeddings(
771
+ inputs_embeds=inputs_embeds,
772
+ multimodal_embeddings=multimodal_embeddings,
773
+ is_multimodal=_require_is_multimodal(is_multimodal),
774
+ )
775
+
776
+ def forward(
777
+ self,
778
+ input_ids: torch.Tensor | None,
779
+ positions: torch.Tensor,
780
+ intermediate_tensors: IntermediateTensors | None = None,
781
+ inputs_embeds: torch.Tensor | None = None,
782
+ **kwargs: object,
783
+ ) -> torch.Tensor | IntermediateTensors:
784
+ if intermediate_tensors is not None:
785
+ inputs_embeds = None
786
+ return self.language_model(
787
+ input_ids=input_ids,
788
+ positions=positions,
789
+ intermediate_tensors=intermediate_tensors,
790
+ inputs_embeds=inputs_embeds,
791
+ )
792
+
793
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
794
+ return self.language_model.compute_logits(hidden_states)
795
+
796
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
797
+ loader = AutoWeightsLoader(self)
798
+ return loader.load_weights(
799
+ weights,
800
+ mapper=self.hf_to_vllm_mapper,
801
+ )
vllm/model_executor/models/registry.py
+4
@@ -447,6 +447,10 @@ _MULTIMODAL_MODELS = {
447
"KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),
448
"KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"),
449
"MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"),
450
+ "MossTranscribeDiarizeForConditionalGeneration": (
451
+ "moss_transcribe_diarize",
452
+ "MossTranscribeDiarizeForConditionalGeneration",
453
+ ),
454
"LightOnOCRForConditionalGeneration": (
455
"lightonocr",
456
"LightOnOCRForConditionalGeneration",
vllm/transformers_utils/config.py
+1
@@ -107,6 +107,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
107
minimax_m3_vl="MiniMaxM3Config",
108
minimax_m3_mtp="MiniMaxM3MTPConfig",
109
moondream3="Moondream3Config",
110
+ moss_transcribe_diarize="MossTranscribeDiarizeConfig",
111
eagle="EAGLEConfig",
112
speculators="SpeculatorsConfig",
113
nemotron="NemotronConfig",
vllm/transformers_utils/configs/__init__.py
+4
@@ -61,6 +61,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
61
"Moondream3Config": "vllm.transformers_utils.configs.moondream3",
62
"Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3",
63
"Moondream3VisionConfig": "vllm.transformers_utils.configs.moondream3",
64
+ "MossTranscribeDiarizeConfig": (
65
+ "vllm.transformers_utils.configs.moss_transcribe_diarize"
66
+ ),
67
"MoonViTConfig": "vllm.transformers_utils.configs.moonvit",
68
"KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
69
"KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
@@ -136,6 +139,7 @@ __all__ = [
139
"Moondream3Config",
140
"Moondream3TextConfig",
141
"Moondream3VisionConfig",
142
+ "MossTranscribeDiarizeConfig",
143
"MoonViTConfig",
144
"KimiLinearConfig",
145
"KimiVLConfig",
vllm/transformers_utils/configs/moss_transcribe_diarize.py
new
+108
@@ -0,0 +1,108 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Any
5
+
6
+from transformers import PretrainedConfig, Qwen3Config
7
+from transformers.models.whisper.configuration_whisper import WhisperConfig
8
+
9
+
10
+class MossTranscribeDiarizeConfig(PretrainedConfig):
11
+ """Configuration for MOSS-Transcribe-Diarize."""
12
+
13
+ model_type = "moss_transcribe_diarize"
14
+ sub_configs = {"text_config": Qwen3Config, "audio_config": WhisperConfig}
15
+ keys_to_ignore_at_inference = ["past_key_values"]
16
+
17
+ def __init__(
18
+ self,
19
+ text_config: dict[str, Any] | Qwen3Config | None = None,
20
+ audio_config: dict[str, Any] | WhisperConfig | None = None,
21
+ audio_token_id: int = 151671,
22
+ audio_merge_size: int = 4,
23
+ adaptor_input_dim: int | None = None,
24
+ tie_word_embeddings: bool = True,
25
+ **kwargs: Any,
26
+ ) -> None:
27
+ text_config_obj: Qwen3Config
28
+ if text_config is None:
29
+ text_config_obj = Qwen3Config(
30
+ vocab_size=151936,
31
+ hidden_size=1024,
32
+ intermediate_size=3072,
33
+ num_hidden_layers=28,
34
+ num_attention_heads=16,
35
+ num_key_value_heads=8,
36
+ head_dim=128,
37
+ max_position_embeddings=40960,
38
+ tie_word_embeddings=tie_word_embeddings,
39
+ rope_theta=1_000_000.0,
40
+ layer_types=["full_attention"] * 28,
41
+ )
42
+ elif isinstance(text_config, dict):
43
+ text_config_obj = Qwen3Config(**text_config)
44
+ else:
45
+ text_config_obj = text_config
46
+
47
+ audio_config_obj: WhisperConfig
48
+ if audio_config is None:
49
+ audio_config_obj = WhisperConfig(
50
+ num_mel_bins=80,
51
+ d_model=1024,
52
+ encoder_layers=24,
53
+ encoder_attention_heads=16,
54
+ encoder_ffn_dim=4096,
55
+ max_source_positions=1500,
56
+ dropout=0.0,
57
+ attention_dropout=0.0,
58
+ activation_dropout=0.0,
59
+ activation_function="gelu",
60
+ encoder_layerdrop=0.0,
61
+ scale_embedding=False,
62
+ )
63
+ elif isinstance(audio_config, dict):
64
+ audio_config_obj = WhisperConfig(**audio_config)
65
+ else:
66
+ audio_config_obj = audio_config
67
+
68
+ text_config_obj.tie_word_embeddings = tie_word_embeddings
69
+ if not getattr(text_config_obj, "layer_types", None):
70
+ text_config_obj.layer_types = [
71
+ "full_attention"
72
+ ] * text_config_obj.num_hidden_layers
73
+
74
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
75
+
76
+ self.text_config = text_config_obj
77
+ self.audio_config = audio_config_obj
78
+ self.audio_token_id = int(audio_token_id)
79
+ self.audio_merge_size = int(audio_merge_size)
80
+ self.adaptor_input_dim = (
81
+ int(adaptor_input_dim)
82
+ if adaptor_input_dim is not None
83
+ else int(audio_config_obj.d_model) * int(audio_merge_size)
84
+ )
85
+
86
+ self.vocab_size = int(text_config_obj.vocab_size)
87
+ self.hidden_size = int(text_config_obj.hidden_size)
88
+ self.intermediate_size = int(text_config_obj.intermediate_size)
89
+ self.num_hidden_layers = int(text_config_obj.num_hidden_layers)
90
+ self.num_attention_heads = int(text_config_obj.num_attention_heads)
91
+ self.num_key_value_heads = int(text_config_obj.num_key_value_heads)
92
+ self.head_dim = int(text_config_obj.head_dim)
93
+ self.hidden_act = text_config_obj.hidden_act
94
+ self.max_position_embeddings = int(text_config_obj.max_position_embeddings)
95
+ self.rms_norm_eps = float(text_config_obj.rms_norm_eps)
96
+ rope_parameters = getattr(text_config_obj, "rope_parameters", None)
97
+ rope_theta = float(getattr(text_config_obj, "rope_theta", 1_000_000.0))
98
+ if rope_parameters is None:
99
+ rope_parameters = {
100
+ "rope_type": "default",
101
+ "rope_theta": rope_theta,
102
+ }
103
+ text_config_obj.rope_parameters = rope_parameters
104
+ self.rope_parameters = rope_parameters
105
+ self.rope_theta = float(rope_parameters.get("rope_theta", rope_theta))
106
+ self.attention_bias = bool(text_config_obj.attention_bias)
107
+ self.attention_dropout = float(text_config_obj.attention_dropout)
108
+ self.is_causal = True