@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9d1aa4dda8

[Frontend] Add diarized_json support for MOSS-Transcribe-Diarize (#48543)

Signed-off-by: Lucas Christian <lucaschgf7@gmail.com>

christian committed Jul 29, 2026 at 20:24 UTC 9d1aa4dda85a58c74b23ff79aa4a20594ee6f359

9 files changed +309 -10

docs/serving/online_serving/speech_to_text.md

+25 -1

@@ -64,7 +64,8 @@ The Transcriptions API supports uploading audio files in various formats includi

64 - `model`: The model to use for transcription (required)

65 - `language`: The language code (e.g., "en", "zh") (optional)

66 - `prompt`: Optional text to guide the transcription style (optional)

67 -- `response_format`: Format of the response ("json", "text") (optional)

67 +- `response_format`: Format of the response ("json", "text", "verbose_json", or

68 + "diarized_json") (optional)

69 - `temperature`: Sampling temperature between 0 and 1 (optional)

71 For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).

@@ -98,6 +99,29 @@ For `verbose_json` response format:

99 ```

100 Currently “verbose_json” response format doesn’t support no_speech_prob.

101

102 +For models with diarization support, `diarized_json` returns OpenAI-compatible

103 +speaker segments. Currently, this is supported by

104 +`OpenMOSS-Team/MOSS-Transcribe-Diarize`.

105 +

106 +```json

107 +{

108 + "task": "transcribe",

109 + "duration": 6.1,

110 + "text": "Hello. Hi, how are you?",

111 + "segments": [

112 + {

113 + "type": "transcript.text.segment",

114 + "id": "segment_0",

115 + "start": 0.0,

116 + "end": 2.8,

117 + "text": "Hello.",

118 + "speaker": "S01"

119 + }

120 + ],

121 + "usage": {"type": "duration", "seconds": 7}

122 +}

123 +```

124 +

125 ### Extra Parameters

126

127 The following [sampling parameters](../../api/README.md#inference-parameters) are supported.

tests/models/multimodal/processing/test_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 vllm.model_executor.models.moss_transcribe_diarize import (

5 + MossTranscribeDiarizeForConditionalGeneration,

6 +)

7 +

8 +

9 +def test_parse_diarized_transcript_preserves_moss_segments():

10 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

11 + "[0.48][S01]Welcome[1.66][12.26][S02]Ready[13.81]"

12 + )

13 +

14 + assert [

15 + (segment.start, segment.end, segment.speaker, segment.text)

16 + for segment in segments

17 + ] == [

18 + (0.48, 1.66, "S01", "Welcome"),

19 + (12.26, 13.81, "S02", "Ready"),

20 + ]

21 +

22 +

23 +def test_parse_diarized_transcript_preserves_overlapping_segments():

24 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

25 + "[0][S01]First speaker[2][1][S02]Second speaker[3]"

26 + )

27 +

28 + assert [

29 + (segment.start, segment.end, segment.speaker, segment.text)

30 + for segment in segments

31 + ] == [

32 + (0.0, 2.0, "S01", "First speaker"),

33 + (1.0, 3.0, "S02", "Second speaker"),

34 + ]

35 +

36 +

37 +def test_parse_diarized_transcript_preserves_numeric_text_markers():

38 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

39 + "[0][S01]The [2024] report is ready.[4]"

40 + )

41 +

42 + assert [segment.text for segment in segments] == ["The [2024] report is ready."]

43 +

44 +

45 +def test_parse_diarized_transcript_ignores_whitespace_between_segments():

46 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

47 + "[0][S01]Hello[1]\n [2][S02]Hi[3]"

48 + )

49 +

50 + assert [(segment.start, segment.end, segment.text) for segment in segments] == [

51 + (0.0, 1.0, "Hello"),

52 + (2.0, 3.0, "Hi"),

53 + ]

54 +

55 +

56 +def test_parse_diarized_transcript_ignores_noise_before_a_segment():

57 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

58 + "noise [bad][0.1][S01]Hello[0.9]"

59 + )

60 +

61 + assert [(segment.start, segment.end, segment.text) for segment in segments] == [

62 + (0.1, 0.9, "Hello"),

63 + ]

64 +

65 +

66 +def test_parse_diarized_transcript_preserves_timestamps_before_the_end():

67 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

68 + "[2][S01]The earlier timestamp is [1] not the end[3]"

69 + )

70 +

71 + assert [segment.text for segment in segments] == [

72 + "The earlier timestamp is [1] not the end",

73 + ]

74 +

75 +

76 +def test_parse_diarized_transcript_skips_empty_segments():

77 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

78 + "[0][S01][1][2][S02]Complete[3]"

79 + )

80 +

81 + assert [(segment.speaker, segment.text) for segment in segments] == [

82 + ("S02", "Complete"),

83 + ]

84 +

85 +

86 +def test_parse_diarized_transcript_fails_closed_for_incomplete_output():

87 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

88 + "[0][S01]Complete[1][2][S02]Incomplete"

89 + )

90 +

91 + assert segments == []

92 +

93 +

94 +def test_parse_diarized_transcript_fails_closed_for_trailing_text():

95 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

96 + "[0][S01]Complete[1] trailing text"

97 + )

98 +

99 + assert segments == []

100 +

101 +

102 +def test_parse_diarized_transcript_preserves_overlong_timestamp_markers():

103 + marker = f"[{'1' * 33}]"

104 + segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(

105 + f"[0][S01]Value {marker}[1]"

106 + )

107 +

108 + assert [segment.text for segment in segments] == [f"Value {marker}"]

vllm/entrypoints/speech_to_text/base/protocol.py

@@ -8,4 +8,5 @@ import torch

9 ## Protocols for Audio

10 AudioResponseFormat: TypeAlias = Literal["json", "text", "srt", "verbose_json", "vtt"]

11 +TranscriptionResponseFormat: TypeAlias = AudioResponseFormat | Literal["diarized_json"]

12 _LONG_INFO = torch.iinfo(torch.long)

vllm/entrypoints/speech_to_text/base/serving.py

+52 -6

@@ -43,6 +43,7 @@ from vllm.utils.async_utils import make_async_with_semaphore, merge_async_iterat

44 from ..transcription.protocol import (

45 TranscriptionResponse,

46 + TranscriptionResponseDiarized,

47 TranscriptionResponseStreamChoice,

48 TranscriptionResponseVerbose,

49 TranscriptionSegment,

@@ -56,7 +57,9 @@ from ..translation.protocol import (

57 TranslationStreamResponse,

58 )

59 -SpeechToTextResponse: TypeAlias = TranscriptionResponse | TranslationResponse

60 +SpeechToTextResponse: TypeAlias = (

61 + TranscriptionResponse | TranslationResponse | TranscriptionResponseDiarized

62 +)

63 SpeechToTextResponseVerbose: TypeAlias = (

64 TranscriptionResponseVerbose | TranslationResponseVerbose

65 )

@@ -70,6 +73,7 @@ ResponseType: TypeAlias = (

73 TranscriptionResponse

74 | TranslationResponse

75 | TranscriptionResponseVerbose

76 + | TranscriptionResponseDiarized

77 | TranslationResponseVerbose

78 )

@@ -448,10 +452,15 @@ class SpeechToTextBaseServing(GenerateBaseServing):

452 if self.engine_client.errored:

453 raise self.engine_client.dead_error

454

451 - if request.response_format not in ["text", "json", "verbose_json"]:

455 + if request.response_format not in [

456 + "text",

457 + "json",

458 + "verbose_json",

459 + "diarized_json",

460 + ]:

461 return self.create_error_response(

462 "Currently only support response_format: "

454 - "`text`, `json` or `verbose_json`"

463 + "`text`, `json`, `verbose_json` or `diarized_json`"

464 )

465

466 if (

@@ -462,9 +471,20 @@ class SpeechToTextBaseServing(GenerateBaseServing):

471 f"Currently do not support verbose_json for {request.model}"

472 )

473

465 - if request.response_format == "verbose_json" and request.stream:

474 + if (

475 + request.response_format == "diarized_json"

476 + and not self.model_cls.supports_diarized_transcription

477 + ):

478 + return self.create_error_response(

479 + f"Currently do not support diarized_json for {request.model}"

480 + )

481 +

482 + if (

483 + request.response_format in {"verbose_json", "diarized_json"}

484 + and request.stream

485 + ):

486 return self.create_error_response(

467 - "verbose_json format doesn't support streaming case"

487 + f"{request.response_format} format doesn't support streaming case"

488 )

489 request_id = f"{self.task_type}-{self._base_request_id(raw_request)}"

490

@@ -634,7 +654,33 @@ class SpeechToTextBaseServing(GenerateBaseServing):

654 # rounded up as per openAI specs

655 "seconds": int(math.ceil(duration_s)),

656 }

637 - if request.response_format != "verbose_json":

657 + if request.response_format == "diarized_json":

658 + diarized_segments = self.model_cls.parse_diarized_transcript(text)

659 + if not diarized_segments:

660 + return self.create_error_response(

661 + "Model output did not contain a valid diarized transcript"

662 + )

663 + final_response = cast(

664 + T,

665 + TranscriptionResponseDiarized(

666 + duration=duration_s,

667 + text=separator.join(

668 + segment.text for segment in diarized_segments

669 + ),

670 + segments=[

671 + {

672 + "id": f"seg_{index}",

673 + "start": segment.start,

674 + "end": segment.end,

675 + "text": segment.text,

676 + "speaker": segment.speaker,

677 + }

678 + for index, segment in enumerate(diarized_segments)

679 + ],

680 + usage=usage,

681 + ),

682 + )

683 + elif request.response_format != "verbose_json":

684 final_response = cast(

685 T, TranscriptionResponse(text=text, usage=usage)

686 )

vllm/entrypoints/speech_to_text/transcription/protocol.py

+24 -3

@@ -27,7 +27,7 @@ from vllm.sampling_params import (

27 )

28 from vllm.utils import random_uuid

30 -from ..base.protocol import _LONG_INFO, AudioResponseFormat

30 +from ..base.protocol import _LONG_INFO, TranscriptionResponseFormat

32 if TYPE_CHECKING:

33 import numpy as np

@@ -88,7 +88,7 @@ class TranscriptionRequest(OpenAIBaseModel):

88 should match the audio language.

89 """

91 - response_format: AudioResponseFormat = Field(default="json")

91 + response_format: TranscriptionResponseFormat = Field(default="json")

92 """

93 The format of the output, in one of these options: `json`, `text`, `srt`,

94 `verbose_json`, or `vtt`.

@@ -400,6 +400,27 @@ class TranscriptionResponseVerbose(OpenAIBaseModel):

400 """Extracted words and their corresponding timestamps."""

401

402

403 +class TranscriptionDiarizedSegment(OpenAIBaseModel):

404 + """A speaker-attributed transcription segment."""

405 +

406 + type: Literal["transcript.text.segment"] = "transcript.text.segment"

407 + id: str

408 + start: float

409 + end: float

410 + text: str

411 + speaker: str

412 +

413 +

414 +class TranscriptionResponseDiarized(OpenAIBaseModel):

415 + """OpenAI-compatible diarized transcription response."""

416 +

417 + task: Literal["transcribe"] = "transcribe"

418 + duration: float

419 + text: str

420 + segments: list[TranscriptionDiarizedSegment]

421 + usage: TranscriptionUsageAudio

422 +

423 +

424 TranscriptionResponseVariant: TypeAlias = (

404 - TranscriptionResponse | TranscriptionResponseVerbose

425 + TranscriptionResponse | TranscriptionResponseVerbose | TranscriptionResponseDiarized

426 )

vllm/entrypoints/speech_to_text/transcription/serving.py

@@ -18,6 +18,7 @@ from ..base.serving import SpeechToTextBaseServing

18 from .protocol import (

19 TranscriptionRequest,

20 TranscriptionResponse,

21 + TranscriptionResponseDiarized,

22 TranscriptionResponseStreamChoice,

23 TranscriptionResponseVerbose,

24 TranscriptionStreamResponse,

@@ -55,6 +56,7 @@ class OpenAIServingTranscription(SpeechToTextBaseServing):

56 ) -> (

57 TranscriptionResponse

58 | TranscriptionResponseVerbose

59 + | TranscriptionResponseDiarized

60 | AsyncGenerator[str, None]

61 | ErrorResponse

62 ):

@@ -70,6 +72,8 @@ class OpenAIServingTranscription(SpeechToTextBaseServing):

72 response_class=(

73 TranscriptionResponseVerbose

74 if request.response_format == "verbose_json"

75 + else TranscriptionResponseDiarized

76 + if request.response_format == "diarized_json"

77 else TranscriptionResponse

78 ),

79 stream_generator_method=self.transcription_stream_generator,

vllm/model_executor/models/interfaces.py

+23

@@ -11,6 +11,7 @@ from collections.abc import (

11 Sequence,

12 )

13 from contextlib import ExitStack, contextmanager, nullcontext

14 +from dataclasses import dataclass

15 from typing import (

16 TYPE_CHECKING,

17 Any,

@@ -71,6 +72,16 @@ The output embeddings must be one of the following formats:

72 """

75 +@dataclass(frozen=True, slots=True)

76 +class DiarizedTranscriptionSegment:

77 + """A timestamped, speaker-attributed segment produced by an ASR model."""

78 +

79 + start: float

80 + end: float

81 + speaker: str

82 + text: str

83 +

84 +

85 class StreamingTranscriptionPostProcessor:

86 """Stateful streaming post-processor for transcription deltas."""

@@ -1133,6 +1144,9 @@ class SupportsTranscription(Protocol):

1144 Enables the segment timestamp option for supported models by setting this to `True`.

1145 """

1146

1147 + supports_diarized_transcription: ClassVar[bool] = False

1148 + """Enables the ``diarized_json`` response format for the model."""

1149 +

1150 supports_explicit_language_detection: ClassVar[bool] = False

1151 """

1152 Transcription models that require an explicit language detection step

@@ -1238,6 +1252,15 @@ class SupportsTranscription(Protocol):

1252 """

1253 return text

1254

1255 + @classmethod

1256 + def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegment]:

1257 + """Parse the model-specific diarized transcript format.

1258 +

1259 + Only models that set ``supports_diarized_transcription`` must override

1260 + this method.

1261 + """

1262 + raise NotImplementedError

1263 +

1264 @classmethod

1265 def get_streaming_post_processor_cls(

1266 cls,

vllm/model_executor/models/moss_transcribe_diarize.py

+50

@@ -14,6 +14,7 @@ import math

14 from collections.abc import Iterable, Mapping, Sequence

15 from typing import Annotated, Any, Literal, TypeAlias

17 +import regex as re

18 import torch

19 from torch import nn

20 from transformers import BatchFeature

@@ -23,6 +24,7 @@ from vllm.config.multimodal import BaseDummyOptions

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,

@@ -35,6 +37,8 @@ from vllm.model_executor.models.utils import (

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,

@@ -77,6 +81,11 @@ DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT = (

81 "并在段末标注结束时间戳，以清晰标明该段语音范围。"

82 )

84 +_MOSS_DIARIZED_HEADER_RE = re.compile(

85 + r"\[(?P<start>[0-9.]{1,32})\]\s*\[(?P<speaker>S[0-9]{1,15})\]"

86 +)

87 +_MOSS_DIARIZED_END_RE = re.compile(r"\[(?P<end>[0-9.]{1,32})\]\s*\Z")

88 +

90 class MossTranscribeDiarizeAudioInputs(TensorSchema):

91 """

@@ -563,6 +572,7 @@ class MossTranscribeDiarizeForConditionalGeneration(

572 supports_transcription = True

573 supports_transcription_only = True

574 supports_segment_timestamp = False

575 + supports_diarized_transcription = True

576 supported_languages = ISO639_1_SUPPORTED_LANGS

577 hf_to_vllm_mapper = WeightsMapper(

578 orig_to_new_prefix={

@@ -631,6 +641,46 @@ class MossTranscribeDiarizeForConditionalGeneration(

641 def post_process_output(cls, text: str) -> str:

642 return text.strip()

643

644 + @classmethod

645 + def parse_diarized_transcript(cls, text: str) -> list[DiarizedTranscriptionSegment]:

646 + """Parse MOSS's canonical ``[start][Sxx]text[end]`` transcript."""

647 + headers: list[tuple[re.Match[str], float, str]] = []

648 + for match in _MOSS_DIARIZED_HEADER_RE.finditer(text):

649 + start = parse_diarized_timestamp(match["start"])

650 + speaker = parse_diarized_speaker(match["speaker"])

651 + if start is not None and speaker is not None:

652 + headers.append((match, start, speaker))

653 +

654 + if not headers:

655 + return []

656 +

657 + segments: list[DiarizedTranscriptionSegment] = []

658 + for index, (header, start, speaker) in enumerate(headers):

659 + next_header_start = (

660 + headers[index + 1][0].start() if index + 1 < len(headers) else len(text)

661 + )

662 + body = text[header.end() : next_header_start]

663 + end_match = _MOSS_DIARIZED_END_RE.search(body)

664 + if end_match is None:

665 + return []

666 +

667 + end = parse_diarized_timestamp(end_match["end"])

668 + if end is None or end < start:

669 + return []

670 +

671 + segment_text = body[: end_match.start()].strip()

672 + if segment_text:

673 + segments.append(

674 + DiarizedTranscriptionSegment(

675 + start=start,

676 + end=end,

677 + speaker=speaker,

678 + text=segment_text,

679 + )

680 + )

681 +

682 + return segments

683 +

684 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

685 super().__init__()

686 self.config = vllm_config.model_config.hf_config

vllm/model_executor/models/utils.py

+22

@@ -1074,3 +1074,25 @@ def scatter_output_slices(

1074 sliced = output[offset : offset + n_tok]

1075 dest[idx] = sliced.clone() if clone else sliced

1076 offset += n_tok

1077 +

1078 +

1079 +def parse_diarized_timestamp(marker: str) -> float | None:

1080 + if (

1081 + not marker

1082 + or not marker.isascii()

1083 + or marker.count(".") > 1

1084 + or not marker.replace(".", "").isdigit()

1085 + ):

1086 + return None

1087 + return float(marker)

1088 +

1089 +

1090 +def parse_diarized_speaker(speaker: str) -> str | None:

1091 + if (

1092 + len(speaker) < 2

1093 + or speaker[0] != "S"

1094 + or not speaker[1:].isascii()

1095 + or not speaker[1:].isdigit()

1096 + ):

1097 + return None

1098 + return speaker