279 lines
8.1 KiB
Python
279 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
from typing import Protocol, cast
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import soundfile as sf
|
|
|
|
import vllm.benchmarks.datasets.datasets as datasets_module
|
|
import vllm.benchmarks.lib.endpoint_request_func as request_func_module
|
|
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
|
|
|
|
pytestmark = pytest.mark.skip_global_cleanup
|
|
|
|
|
|
class _ReadableBinary(Protocol):
|
|
def read(self, size: int = -1) -> bytes: ...
|
|
|
|
|
|
class _TokenizedPrompt:
|
|
def __init__(self, prompt: str) -> None:
|
|
self.input_ids = prompt.split()
|
|
|
|
|
|
class _Tokenizer:
|
|
def __init__(self, name_or_path: str = "openai/whisper-large-v3") -> None:
|
|
self.name_or_path = name_or_path
|
|
|
|
def __call__(self, prompt: str) -> _TokenizedPrompt:
|
|
return _TokenizedPrompt(prompt)
|
|
|
|
|
|
class CohereAsrTokenizer(_Tokenizer):
|
|
def __init__(self, name_or_path: str = "/models/cohere-transcribe") -> None:
|
|
super().__init__(name_or_path)
|
|
|
|
|
|
class _CohereNameOnlyTokenizer(_Tokenizer):
|
|
def __init__(self) -> None:
|
|
super().__init__("cohere/some-local-checkpoint")
|
|
|
|
|
|
def _write_wav(path: Path, duration_s: float = 0.1, sample_rate: int = 16_000) -> None:
|
|
num_samples = int(duration_s * sample_rate)
|
|
sf.write(path, np.zeros(num_samples, dtype=np.float32), sample_rate)
|
|
|
|
|
|
class _FakeFormData:
|
|
def __init__(self) -> None:
|
|
self.fields: list[tuple[str, object, dict[str, str]]] = []
|
|
|
|
def add_field(self, name: str, value: object, **kwargs: str) -> None:
|
|
self.fields.append((name, value, kwargs))
|
|
|
|
|
|
class _FakeContent:
|
|
async def iter_any(self):
|
|
yield b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'
|
|
yield b'data: {"usage":{"completion_tokens":1}}\n\n'
|
|
yield b"data: [DONE]\n\n"
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self) -> None:
|
|
self.status = 200
|
|
self.reason = "OK"
|
|
self.content = _FakeContent()
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
|
|
class _FakeSession:
|
|
def __init__(self) -> None:
|
|
self.uploaded_bytes: bytes | None = None
|
|
self.upload_filename: str | None = None
|
|
self.fields: list[tuple[str, object, dict[str, str]]] | None = None
|
|
|
|
def post(self, *, url: str, data: _FakeFormData, headers: dict[str, str]):
|
|
del url, headers
|
|
self.fields = list(data.fields)
|
|
_, file_obj, file_kwargs = self.fields[0]
|
|
file_obj = cast(_ReadableBinary, file_obj)
|
|
self.uploaded_bytes = file_obj.read()
|
|
self.upload_filename = file_kwargs.get("filename")
|
|
return _FakeResponse()
|
|
|
|
|
|
def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None:
|
|
audio_path = tmp_path / "earnings.wav"
|
|
_write_wav(audio_path, duration_s=0.1)
|
|
|
|
dataset = object.__new__(datasets_module.ASRDataset)
|
|
dataset.data = [
|
|
{
|
|
"audio": {
|
|
"path": str(audio_path),
|
|
"bytes": None,
|
|
},
|
|
"text": "quarterly earnings call",
|
|
}
|
|
]
|
|
|
|
samples = dataset.sample(
|
|
tokenizer=_Tokenizer(),
|
|
num_requests=1,
|
|
output_len=32,
|
|
asr_min_audio_len_sec=0.0,
|
|
asr_max_audio_len_sec=1.0,
|
|
)
|
|
|
|
assert len(samples) == 1
|
|
assert samples[0].multi_modal_data == {"audio_path": str(audio_path)}
|
|
assert (
|
|
samples[0].prompt == "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("has_filepath", [True, False])
|
|
def test_asr_dataset_sample_handles_embedded_audio_bytes(
|
|
tmp_path: Path, has_filepath: bool
|
|
) -> None:
|
|
audio_path = tmp_path / "earnings.wav"
|
|
_write_wav(audio_path, duration_s=0.1)
|
|
|
|
test_path = None
|
|
if has_filepath:
|
|
test_path = audio_path
|
|
|
|
dataset = object.__new__(datasets_module.ASRDataset)
|
|
dataset.data = [
|
|
{
|
|
"audio": {
|
|
"path": test_path,
|
|
"bytes": audio_path.read_bytes(),
|
|
},
|
|
"text": "quarterly earnings call",
|
|
}
|
|
]
|
|
|
|
samples = dataset.sample(
|
|
tokenizer=_Tokenizer(),
|
|
num_requests=1,
|
|
output_len=32,
|
|
asr_min_audio_len_sec=0.0,
|
|
asr_max_audio_len_sec=1.0,
|
|
)
|
|
|
|
assert len(samples) == 1
|
|
assert isinstance(samples[0].multi_modal_data, dict)
|
|
audio, sample_rate = samples[0].multi_modal_data["audio"]
|
|
assert sample_rate == 16_000
|
|
assert isinstance(audio, np.ndarray)
|
|
assert audio.size > 0
|
|
|
|
|
|
def test_async_request_openai_audio_handles_local_audio_paths(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
audio_path = tmp_path / "earnings.wav"
|
|
_write_wav(audio_path, duration_s=0.25)
|
|
|
|
monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData)
|
|
session = _FakeSession()
|
|
request_input = RequestFuncInput(
|
|
prompt="",
|
|
api_url="http://localhost:8000/v1/audio/transcriptions",
|
|
prompt_len=1,
|
|
output_len=32,
|
|
model="openai/whisper-large-v3",
|
|
multi_modal_content={"audio_path": str(audio_path)},
|
|
)
|
|
|
|
output = asyncio.run(
|
|
request_func_module.async_request_openai_audio(request_input, session)
|
|
)
|
|
|
|
assert session.upload_filename == audio_path.name
|
|
assert session.uploaded_bytes == audio_path.read_bytes()
|
|
assert output.success is True
|
|
assert output.generated_text == "hello"
|
|
assert output.output_tokens == 1
|
|
assert output.input_audio_duration == pytest.approx(0.25, abs=1e-2)
|
|
|
|
|
|
def test_async_request_openai_audio_handles_decoded_audio_arrays(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData)
|
|
session = _FakeSession()
|
|
request_input = RequestFuncInput(
|
|
prompt="",
|
|
api_url="http://localhost:8000/v1/audio/transcriptions",
|
|
prompt_len=1,
|
|
output_len=32,
|
|
model="openai/whisper-large-v3",
|
|
multi_modal_content={
|
|
"audio": (np.zeros(1_600, dtype=np.float32), 16_000),
|
|
},
|
|
)
|
|
|
|
output = asyncio.run(
|
|
request_func_module.async_request_openai_audio(request_input, session)
|
|
)
|
|
|
|
assert session.upload_filename == "audio.wav"
|
|
assert session.uploaded_bytes is not None
|
|
assert output.success is True
|
|
assert output.generated_text == "hello"
|
|
|
|
|
|
_COHERE_ASR_PROMPT = (
|
|
"<|startofcontext|><|startoftranscript|>"
|
|
"<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>"
|
|
"<|notimestamp|><|nodiarize|>"
|
|
)
|
|
|
|
|
|
def _make_asr_dataset(tmp_path: Path) -> datasets_module.ASRDataset:
|
|
audio_path = tmp_path / "sample.wav"
|
|
_write_wav(audio_path, duration_s=0.1)
|
|
dataset = object.__new__(datasets_module.ASRDataset)
|
|
dataset.data = [
|
|
{
|
|
"audio": {"path": str(audio_path), "bytes": None},
|
|
"text": "hello world",
|
|
}
|
|
]
|
|
return dataset
|
|
|
|
|
|
def test_asr_dataset_cohere_class_name_gets_decoder_prompt(tmp_path: Path) -> None:
|
|
dataset = _make_asr_dataset(tmp_path)
|
|
samples = dataset.sample(
|
|
tokenizer=CohereAsrTokenizer(),
|
|
num_requests=1,
|
|
output_len=32,
|
|
asr_min_audio_len_sec=0.0,
|
|
asr_max_audio_len_sec=1.0,
|
|
)
|
|
assert len(samples) == 1
|
|
assert samples[0].prompt == _COHERE_ASR_PROMPT
|
|
|
|
|
|
def test_asr_dataset_cohere_name_or_path_fallback_gets_decoder_prompt(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset = _make_asr_dataset(tmp_path)
|
|
samples = dataset.sample(
|
|
tokenizer=_CohereNameOnlyTokenizer(),
|
|
num_requests=1,
|
|
output_len=32,
|
|
asr_min_audio_len_sec=0.0,
|
|
asr_max_audio_len_sec=1.0,
|
|
)
|
|
assert len(samples) == 1
|
|
assert samples[0].prompt == _COHERE_ASR_PROMPT
|
|
|
|
|
|
def test_asr_dataset_unknown_tokenizer_gets_empty_prompt(tmp_path: Path) -> None:
|
|
dataset = _make_asr_dataset(tmp_path)
|
|
samples = dataset.sample(
|
|
tokenizer=_Tokenizer(name_or_path="some-other/asr-model"),
|
|
num_requests=1,
|
|
output_len=32,
|
|
asr_min_audio_len_sec=0.0,
|
|
asr_max_audio_len_sec=1.0,
|
|
)
|
|
assert len(samples) == 1
|
|
assert samples[0].prompt == ""
|