[Bugfix][Model] Fix CohereASR streaming audio-token estimate (unit + subsampling) (#53829)
Signed-off-by: hungh <hungh@nvidia.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Hung Hoang (黄兴） committed
Sep 2, 2026 at 22:55 UTC
3b45d053b4bbc61ce437f00891b52ce5ddde7c5a
2 files changed
+43
-3
tests/models/multimodal/test_cohere_asr.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+
8
+from vllm.config import SpeechToTextConfig
9
+from vllm.model_executor.models.cohere_asr import CohereAsrForConditionalGeneration
10
+
11
+
12
+@pytest.mark.parametrize(
13
+ ("audio_duration_s", "expected_tokens"),
14
+ # window_stride=0.01s @ 16kHz -> 160-sample hop; subsampling_factor=8.
15
+ # frames = floor(duration * 16000 / 160); tokens = ceil(frames / 8).
16
+ [(1.0, 13), (10.0, 125), (30.0, 375)],
17
+)
18
+def test_get_num_audio_tokens_streaming_estimate(audio_duration_s, expected_tokens):
19
+ """The duration-based estimate must convert ``window_stride`` (seconds) to a
20
+ sample hop and divide by the encoder subsampling factor. Values are pinned to
21
+ concrete numbers so a rounding regression is actually caught."""
22
+ model_config = SimpleNamespace(
23
+ hf_config=SimpleNamespace(
24
+ preprocessor={"window_stride": 0.01, "sample_rate": 16000},
25
+ encoder={"subsampling_factor": 8},
26
+ )
27
+ )
28
+ stt_config = SpeechToTextConfig(sample_rate=16000)
29
+
30
+ got = CohereAsrForConditionalGeneration.get_num_audio_tokens(
31
+ audio_duration_s, stt_config, model_config
32
+ )
33
+
34
+ assert got == expected_tokens
vllm/model_executor/models/cohere_asr.py
+9
-3
@@ -2177,9 +2177,15 @@ class CohereAsrForConditionalGeneration(
2177
stt_config: SpeechToTextConfig,
2178
model_config: ModelConfig,
2179
) -> int | None:
2180
- hop_length = model_config.hf_config.preprocessor.get("window_stride")
2181
- assert hop_length is not None
2182
- return math.ceil(audio_duration_s * stt_config.sample_rate / hop_length)
2180
+ hf_config = model_config.hf_config
2181
+ preprocessor = hf_config.preprocessor
2182
+ sample_rate = preprocessor["sample_rate"]
2183
+ window_stride = preprocessor["window_stride"]
2184
+ hop_length = int(window_stride * sample_rate)
2185
+ # Floor-divide to match get_seq_len.
2186
+ num_frames = int(audio_duration_s * sample_rate) // hop_length
2187
+ subsampling_factor = hf_config.encoder["subsampling_factor"]
2188
+ return math.ceil(num_frames / subsampling_factor)
2189
2190
def get_num_encoder_cross_attn_tokens(self, num_encoder_input_tokens: int) -> int:
2191
return self.model.encoder.get_num_encoder_cross_attn_tokens(