| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | |
| 5 | import asyncio |
| 6 | from pathlib import Path |
| 7 | from typing import Protocol, cast |
| 8 | |
| 9 | import numpy as np |
| 10 | import pytest |
| 11 | import soundfile as sf |
| 12 | |
| 13 | import vllm.benchmarks.datasets.datasets as datasets_module |
| 14 | import vllm.benchmarks.lib.endpoint_request_func as request_func_module |
| 15 | from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput |
| 16 | from vllm.tokenizers import TokenizerLike |
| 17 | |
| 18 | pytestmark = pytest.mark.skip_global_cleanup |
| 19 | |
| 20 | |
| 21 | class _ReadableBinary(Protocol): |
| 22 | def read(self, size: int = -1) -> bytes: ... |
| 23 | |
| 24 | |
| 25 | class _TokenizedPrompt: |
| 26 | def __init__(self, prompt: str) -> None: |
| 27 | self.input_ids = prompt.split() |
| 28 | |
| 29 | |
| 30 | class _Tokenizer: |
| 31 | def __init__(self, name_or_path: str = "openai/whisper-large-v3") -> None: |
| 32 | self.name_or_path = name_or_path |
| 33 | |
| 34 | def __call__(self, prompt: str) -> _TokenizedPrompt: |
| 35 | return _TokenizedPrompt(prompt) |
| 36 | |
| 37 | |
| 38 | class CohereAsrTokenizer(_Tokenizer): |
| 39 | def __init__(self, name_or_path: str = "/models/cohere-transcribe") -> None: |
| 40 | super().__init__(name_or_path) |
| 41 | |
| 42 | |
| 43 | class _CohereNameOnlyTokenizer(_Tokenizer): |
| 44 | def __init__(self) -> None: |
| 45 | super().__init__("cohere/some-local-checkpoint") |
| 46 | |
| 47 | |
| 48 | def _as_tokenizer(tokenizer: _Tokenizer) -> TokenizerLike: |
| 49 | return cast(TokenizerLike, tokenizer) |
| 50 | |
| 51 | |
| 52 | def _write_wav(path: Path, duration_s: float = 0.1, sample_rate: int = 16_000) -> None: |
| 53 | num_samples = int(duration_s * sample_rate) |
| 54 | sf.write(path, np.zeros(num_samples, dtype=np.float32), sample_rate) |
| 55 | |
| 56 | |
| 57 | class _FakeFormData: |
| 58 | def __init__(self) -> None: |
| 59 | self.fields: list[tuple[str, object, dict[str, str]]] = [] |
| 60 | |
| 61 | def add_field(self, name: str, value: object, **kwargs: str) -> None: |
| 62 | self.fields.append((name, value, kwargs)) |
| 63 | |
| 64 | |
| 65 | class _FakeContent: |
| 66 | async def iter_any(self): |
| 67 | yield b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' |
| 68 | yield b'data: {"usage":{"completion_tokens":1}}\n\n' |
| 69 | yield b"data: [DONE]\n\n" |
| 70 | |
| 71 | |
| 72 | class _FakeResponse: |
| 73 | def __init__(self) -> None: |
| 74 | self.status = 200 |
| 75 | self.reason = "OK" |
| 76 | self.content = _FakeContent() |
| 77 | |
| 78 | async def __aenter__(self): |
| 79 | return self |
| 80 | |
| 81 | async def __aexit__(self, exc_type, exc, tb): |
| 82 | return False |
| 83 | |
| 84 | |
| 85 | class _FakeSession: |
| 86 | def __init__(self) -> None: |
| 87 | self.uploaded_bytes: bytes | None = None |
| 88 | self.upload_filename: str | None = None |
| 89 | self.fields: list[tuple[str, object, dict[str, str]]] | None = None |
| 90 | |
| 91 | def post(self, *, url: str, data: _FakeFormData, headers: dict[str, str]): |
| 92 | del url, headers |
| 93 | self.fields = list(data.fields) |
| 94 | _, file_obj, file_kwargs = self.fields[0] |
| 95 | file_obj = cast(_ReadableBinary, file_obj) |
| 96 | self.uploaded_bytes = file_obj.read() |
| 97 | self.upload_filename = file_kwargs.get("filename") |
| 98 | return _FakeResponse() |
| 99 | |
| 100 | |
| 101 | def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None: |
| 102 | audio_path = tmp_path / "earnings.wav" |
| 103 | _write_wav(audio_path, duration_s=0.1) |
| 104 | |
| 105 | dataset = object.__new__(datasets_module.ASRDataset) |
| 106 | dataset.data = [ |
| 107 | { |
| 108 | "audio": { |
| 109 | "path": str(audio_path), |
| 110 | "bytes": None, |
| 111 | }, |
| 112 | "text": "quarterly earnings call", |
| 113 | } |
| 114 | ] |
| 115 | |
| 116 | samples = dataset.sample( |
| 117 | tokenizer=_as_tokenizer(_Tokenizer()), |
| 118 | num_requests=1, |
| 119 | output_len=32, |
| 120 | asr_min_audio_len_sec=0.0, |
| 121 | asr_max_audio_len_sec=1.0, |
| 122 | ) |
| 123 | |
| 124 | assert len(samples) == 1 |
| 125 | assert samples[0].multi_modal_data == {"audio_path": str(audio_path)} |
| 126 | assert ( |
| 127 | samples[0].prompt == "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | @pytest.mark.parametrize("has_filepath", [True, False]) |
| 132 | def test_asr_dataset_sample_handles_embedded_audio_bytes( |
| 133 | tmp_path: Path, has_filepath: bool |
| 134 | ) -> None: |
| 135 | audio_path = tmp_path / "earnings.wav" |
| 136 | _write_wav(audio_path, duration_s=0.1) |
| 137 | |
| 138 | test_path = None |
| 139 | if has_filepath: |
| 140 | test_path = audio_path |
| 141 | |
| 142 | dataset = object.__new__(datasets_module.ASRDataset) |
| 143 | dataset.data = [ |
| 144 | { |
| 145 | "audio": { |
| 146 | "path": test_path, |
| 147 | "bytes": audio_path.read_bytes(), |
| 148 | }, |
| 149 | "text": "quarterly earnings call", |
| 150 | } |
| 151 | ] |
| 152 | |
| 153 | samples = dataset.sample( |
| 154 | tokenizer=_as_tokenizer(_Tokenizer()), |
| 155 | num_requests=1, |
| 156 | output_len=32, |
| 157 | asr_min_audio_len_sec=0.0, |
| 158 | asr_max_audio_len_sec=1.0, |
| 159 | ) |
| 160 | |
| 161 | assert len(samples) == 1 |
| 162 | assert isinstance(samples[0].multi_modal_data, dict) |
| 163 | audio, sample_rate = samples[0].multi_modal_data["audio"] |
| 164 | assert sample_rate == 16_000 |
| 165 | assert isinstance(audio, np.ndarray) |
| 166 | assert audio.size > 0 |
| 167 | |
| 168 | |
| 169 | def test_async_request_openai_audio_handles_local_audio_paths( |
| 170 | monkeypatch: pytest.MonkeyPatch, |
| 171 | tmp_path: Path, |
| 172 | ) -> None: |
| 173 | audio_path = tmp_path / "earnings.wav" |
| 174 | _write_wav(audio_path, duration_s=0.25) |
| 175 | |
| 176 | monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData) |
| 177 | session = _FakeSession() |
| 178 | request_input = RequestFuncInput( |
| 179 | prompt="", |
| 180 | api_url="http://localhost:8000/v1/audio/transcriptions", |
| 181 | prompt_len=1, |
| 182 | output_len=32, |
| 183 | model="openai/whisper-large-v3", |
| 184 | multi_modal_content={"audio_path": str(audio_path)}, |
| 185 | ) |
| 186 | |
| 187 | output = asyncio.run( |
| 188 | request_func_module.async_request_openai_audio(request_input, session) |
| 189 | ) |
| 190 | |
| 191 | assert session.upload_filename == audio_path.name |
| 192 | assert session.uploaded_bytes == audio_path.read_bytes() |
| 193 | assert output.success is True |
| 194 | assert output.generated_text == "hello" |
| 195 | assert output.output_tokens == 1 |
| 196 | assert output.input_audio_duration == pytest.approx(0.25, abs=1e-2) |
| 197 | |
| 198 | |
| 199 | def test_async_request_openai_audio_handles_decoded_audio_arrays( |
| 200 | monkeypatch: pytest.MonkeyPatch, |
| 201 | ) -> None: |
| 202 | monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData) |
| 203 | session = _FakeSession() |
| 204 | request_input = RequestFuncInput( |
| 205 | prompt="", |
| 206 | api_url="http://localhost:8000/v1/audio/transcriptions", |
| 207 | prompt_len=1, |
| 208 | output_len=32, |
| 209 | model="openai/whisper-large-v3", |
| 210 | multi_modal_content={ |
| 211 | "audio": (np.zeros(1_600, dtype=np.float32), 16_000), |
| 212 | }, |
| 213 | ) |
| 214 | |
| 215 | output = asyncio.run( |
| 216 | request_func_module.async_request_openai_audio(request_input, session) |
| 217 | ) |
| 218 | |
| 219 | assert session.upload_filename == "audio.wav" |
| 220 | assert session.uploaded_bytes is not None |
| 221 | assert output.success is True |
| 222 | assert output.generated_text == "hello" |
| 223 | |
| 224 | |
| 225 | _COHERE_ASR_PROMPT = ( |
| 226 | "<|startofcontext|><|startoftranscript|>" |
| 227 | "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>" |
| 228 | "<|notimestamp|><|nodiarize|>" |
| 229 | ) |
| 230 | |
| 231 | |
| 232 | def _make_asr_dataset(tmp_path: Path) -> datasets_module.ASRDataset: |
| 233 | audio_path = tmp_path / "sample.wav" |
| 234 | _write_wav(audio_path, duration_s=0.1) |
| 235 | dataset = object.__new__(datasets_module.ASRDataset) |
| 236 | dataset.data = [ |
| 237 | { |
| 238 | "audio": {"path": str(audio_path), "bytes": None}, |
| 239 | "text": "hello world", |
| 240 | } |
| 241 | ] |
| 242 | return dataset |
| 243 | |
| 244 | |
| 245 | def test_asr_dataset_cohere_class_name_gets_decoder_prompt(tmp_path: Path) -> None: |
| 246 | dataset = _make_asr_dataset(tmp_path) |
| 247 | samples = dataset.sample( |
| 248 | tokenizer=_as_tokenizer(CohereAsrTokenizer()), |
| 249 | num_requests=1, |
| 250 | output_len=32, |
| 251 | asr_min_audio_len_sec=0.0, |
| 252 | asr_max_audio_len_sec=1.0, |
| 253 | ) |
| 254 | assert len(samples) == 1 |
| 255 | assert samples[0].prompt == _COHERE_ASR_PROMPT |
| 256 | |
| 257 | |
| 258 | def test_asr_dataset_cohere_name_or_path_fallback_gets_decoder_prompt( |
| 259 | tmp_path: Path, |
| 260 | ) -> None: |
| 261 | dataset = _make_asr_dataset(tmp_path) |
| 262 | samples = dataset.sample( |
| 263 | tokenizer=_as_tokenizer(_CohereNameOnlyTokenizer()), |
| 264 | num_requests=1, |
| 265 | output_len=32, |
| 266 | asr_min_audio_len_sec=0.0, |
| 267 | asr_max_audio_len_sec=1.0, |
| 268 | ) |
| 269 | assert len(samples) == 1 |
| 270 | assert samples[0].prompt == _COHERE_ASR_PROMPT |
| 271 | |
| 272 | |
| 273 | def test_asr_dataset_unknown_tokenizer_gets_empty_prompt(tmp_path: Path) -> None: |
| 274 | dataset = _make_asr_dataset(tmp_path) |
| 275 | samples = dataset.sample( |
| 276 | tokenizer=_as_tokenizer(_Tokenizer(name_or_path="some-other/asr-model")), |
| 277 | num_requests=1, |
| 278 | output_len=32, |
| 279 | asr_min_audio_len_sec=0.0, |
| 280 | asr_max_audio_len_sec=1.0, |
| 281 | ) |
| 282 | assert len(samples) == 1 |
| 283 | assert samples[0].prompt == "" |