iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Implement custom dataset class for ASR benchmarking (#41576)

Browse Source 
Signed-off-by: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Yasmin Moslem
2026-05-12 12:17:58 +08:00
committed by GitHub
co-authored by gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
parent 630492da30
commit 28ee78af54
3 changed files with 216 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/benchmarking/cli.md
+60 -6
View File
@@ -39,7 +39,8 @@ th {
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
| Custom Audio | ✅ | ✅ | Local file: `audio_data.jsonl` |
| Custom Image | ✅ | ✅ | Local file: `image_data.jsonl` |
Legend:
@@ -142,7 +143,7 @@ Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_da
#### Custom Dataset
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. At inference time, use the option `--dataset-name custom`. Your data needs to be in the `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
```json
{"prompt": "What is the capital of India?"}
@@ -173,9 +174,62 @@ vllm bench serve --port 9001 --save-result --save-detailed \
You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`.
#### Custom multimodal dataset
#### Custom Audio Dataset
If the multimodal dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomMMDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" and "image_files" field per entry, e.g., `mm_data.jsonl`:
If the audio dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomAudioDataset`. At inference time, use the option `--dataset-name custom_audio`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "audio" fields per entry, e.g., `audio_data.jsonl`:
```json
{"prompt": "What does this audio say?", "audio": "/path/to/audio_1.wav"}
{"prompt": "Transcribe the audio.", "audio": "/path/to/audio_2.wav"}
```
- **Supported models:** The `CustomAudioDataset` class supports two types of audio models: ASR models (e.g. Whisper) which do not require a "prompt" field; and multimodal audio-text chat models (e.g. Qwen2-Audio). Since these model types require different arguments at inference, we are giving two examples.
- **Example 1: Whisper**
Whisper is a dedicated ASR encoder-decoder model, so it uses `--backend openai-audio` and `--endpoint /v1/audio/transcriptions`.
```bash
# start server
vllm serve openai/whisper-tiny
```
```bash
vllm bench serve \
--model openai/whisper-tiny \
--backend openai-audio \
--endpoint /v1/audio/transcriptions \
--dataset-name custom_audio \
--dataset-path audio_data.jsonl \
--no-oversample \
--custom-output-len 256 \
--save-result \
--save-detailed \
--result-filename whisper_bench.json
```
- **Example 2: Qwen2-Audio**
Qwen2-Audio is a multimodal chat model that can do ASR and speech analysis, so it uses `--backend openai-chat`, and `--endpoint /v1/chat/completions`. It also requires `--enable-multimodal-chat` to enable multimodal chat transformation.
```bash
vllm bench serve \
--model Qwen/Qwen2-Audio-7B-Instruct \
--backend openai-chat \
--endpoint /v1/chat/completions \
--dataset-name custom_audio \
--dataset-path audio_data.jsonl \
--no-oversample \
--custom-output-len 256 \
--enable-multimodal-chat \
--save-result \
--save-detailed \
--result-filename qwen_bench.json
```
#### Custom Image Dataset
If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
```json
{"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
@@ -193,8 +247,8 @@ vllm bench serve--save-result --save-detailed \
--backend openai-chat \
--model Qwen/Qwen2-VL-7B-Instruct \
--endpoint /v1/chat/completions \
--dataset-name custom_mm \
--dataset-path <path-to-your-mm-data-jsonl> \
--dataset-name custom_image \
--dataset-path <path-to-your-image-data-jsonl> \
--allowed-local-media-path /path/to/image/folder
```
vllm/benchmarks/datasets/__init__.py
+6 -2
View File
@@ -9,8 +9,9 @@ from vllm.benchmarks.datasets.datasets import (
BlazeditDataset,
BurstGPTDataset,
ConversationDataset,
CustomAudioDataset,
CustomDataset,
CustomMMDataset,
CustomImageDataset,
HuggingFaceDataset,
InstructCoderDataset,
MLPerfDataset,
@@ -36,6 +37,7 @@ from vllm.benchmarks.datasets.datasets import (
is_valid_sequence,
lora_path_on_disk,
lora_tokenizer_cache,
process_audio,
process_image,
process_video,
zeta_prompt,
@@ -51,7 +53,8 @@ __all__ = [
"BurstGPTDataset",
"ConversationDataset",
"CustomDataset",
"CustomMMDataset",
"CustomAudioDataset",
"CustomImageDataset",
"HuggingFaceDataset",
"InstructCoderDataset",
"MLPerfDataset",
@@ -77,6 +80,7 @@ __all__ = [
"is_valid_sequence",
"lora_path_on_disk",
"lora_tokenizer_cache",
"process_audio",
"process_image",
"process_video",
"RangeRatio",
vllm/benchmarks/datasets/datasets.py
+150 -4
View File
@@ -60,6 +60,11 @@ try:
except ImportError:
pd = PlaceholderModule("pandas")
try:
import soundfile as sf
except ImportError:
sf = PlaceholderModule("soundfile")
logger = logging.getLogger(__name__)
@@ -441,6 +446,27 @@ def process_video(video: Any) -> Mapping[str, Any]:
)
def process_audio(audio: Any) -> tuple:
"""
Process a single audio input and return a (array, sample_rate) tuple.
Supports:
1. String: treated as a file path, loaded with soundfile.
2. Dict with 'array' and 'sampling_rate' keys: HuggingFace audio format.
3. Tuple (array, sr): passed through directly.
"""
if isinstance(audio, str):
return sf.read(audio)
if isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
return audio["array"], audio["sampling_rate"]
if isinstance(audio, tuple) and len(audio) == 2:
return audio
raise ValueError(
f"Invalid audio input {audio}. Must be a file path string, "
"a dict with 'array' and 'sampling_rate', or a (array, sr) tuple."
)
def gen_prompt_decode_to_target_len(
tokenizer: TokenizerLike,
token_sequence: list[int],
@@ -1399,6 +1425,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
"random-rerank",
"hf",
"custom",
"custom_audio",
"custom_image",
"custom_mm",
"prefix_repetition",
"spec_bench",
@@ -1816,8 +1844,28 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
no_oversample=args.no_oversample,
)
elif args.dataset_name == "custom_mm":
dataset = CustomMMDataset(
elif args.dataset_name in ("custom_image", "custom_mm"):
if args.dataset_name == "custom_mm":
logger.warning(
"Dataset name 'custom_mm' is deprecated and will be removed in v0.24. "
"Use '--dataset-name custom_image' instead."
)
dataset = CustomImageDataset(
dataset_path=args.dataset_path,
disable_shuffle=args.disable_shuffle,
random_seed=args.seed,
)
input_requests = dataset.sample(
num_requests=args.num_prompts,
tokenizer=tokenizer,
output_len=args.custom_output_len,
enable_multimodal_chat=args.enable_multimodal_chat,
request_id_prefix=args.request_id_prefix,
no_oversample=args.no_oversample,
)
elif args.dataset_name == "custom_audio":
dataset = CustomAudioDataset(
dataset_path=args.dataset_path,
disable_shuffle=args.disable_shuffle,
random_seed=args.seed,
@@ -2249,9 +2297,9 @@ class CustomDataset(BenchmarkDataset):
return sampled_requests
class CustomMMDataset(CustomDataset):
class CustomImageDataset(CustomDataset):
"""
Implements the Custom MultiModal dataset. Loads data from a JSONL file and generates
Implements the Custom image dataset. Loads data from a JSONL file and generates
sample requests based on conversation turns. E.g.,
```
{
@@ -2328,6 +2376,104 @@ class CustomMMDataset(CustomDataset):
return sampled_requests
class CustomAudioDataset(CustomDataset):
"""
Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g.,
{"prompt": "Transcribe the audio.", "audio": "/path/to/audio.wav"}
Supports both:
- Dedicated ASR models (e.g. Whisper) via openai-audio & /v1/audio/transcriptions
- Chat-based audio models (e.g. Qwen2-Audio) via openai-chat & /v1/chat/completions
"""
IS_MULTIMODAL = True
def sample(
self,
tokenizer: TokenizerLike,
num_requests: int,
output_len: int | None = None,
request_id_prefix: str = "",
no_oversample: bool = False,
skip_chat_template: bool = False,
enable_multimodal_chat: bool = False,
**kwargs,
) -> list[SampleRequest]:
self.num_available_samples = len(self.data)
if num_requests <= 0:
num_requests = self.num_available_samples
sampled_requests = []
for i, item in enumerate(self.data):
if len(sampled_requests) >= num_requests:
break
prompt = item.get("prompt", "")
if tokenizer is None:
prompt_len = 1
new_output_len = output_len if output_len not in (None, -1) else 256
mm_content = None
else:
use_chat_template = (
not skip_chat_template
and hasattr(tokenizer, "chat_template")
and tokenizer.chat_template is not None
)
if enable_multimodal_chat:
# Chat-based audio models (e.g., Qwen2-Audio):
# encode audio as base64; serve.py assembles the chat message
# as: {"role": "user", "content": [
# {"type": "text", "text": prompt},
# {"type": "input_audio", "input_audio": {...}}
# ]}
y, sr = process_audio(item["audio"])
buf = io.BytesIO()
sf.write(buf, y, sr, format="WAV")
audio_base64 = base64.b64encode(buf.getvalue()).decode("utf-8")
mm_content = {
"type": "input_audio",
"input_audio": {
"data": audio_base64,
"format": "wav",
},
}
# prompt stays as plain string; serve.py handles wrapping
else:
# Whisper-style models: load audio array locally
y, sr = process_audio(item["audio"])
mm_content = {"audio": (y, sr)}
if use_chat_template:
# ASR models with a chat template but not multimodal chat
prompt = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
add_generation_prompt=True,
tokenize=False,
)
# else: plain prompt for Whisper-style models
prompt_len = (
len(tokenizer(prompt).input_ids) if isinstance(prompt, str) else 1
)
new_output_len = output_len
if output_len is None or output_len == -1:
if "output_tokens" not in item:
raise ValueError(
"If no output length is provided the "
"custom dataset must contain an 'output_tokens' field."
)
new_output_len = int(item["output_tokens"])
sampled_requests.append(
SampleRequest(
prompt=prompt,
prompt_len=prompt_len,
expected_output_len=new_output_len,
multi_modal_data=mm_content,
request_id=request_id_prefix + str(i),
)
)
self.maybe_oversample_requests(
sampled_requests, num_requests, request_id_prefix, no_oversample
)
return sampled_requests
# -----------------------------------------------------------------------------
# Spec Bench Dataset Implementation
# -----------------------------------------------------------------------------