A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 21:30 UTC
Implement custom dataset class for ASR benchmarking (#41576)
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>
Yasmin Moslem committed
May 12, 2026 at 05:17 UTC
28ee78af543c563a2fbf78829a7688120e4e4eb5
3 files changed
+216
-12
docs/benchmarking/cli.md
+60
-6
@@ -39,7 +39,8 @@ th {
39
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
40
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
41
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
42
-| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
42
+| Custom Audio | ✅ | ✅ | Local file: `audio_data.jsonl` |
43
+| Custom Image | ✅ | ✅ | Local file: `image_data.jsonl` |
44
45
Legend:
46
@@ -142,7 +143,7 @@ Example output: ![Dataset Statistics](../assets/contributing/vllm_bench_serve_da
143
144
#### Custom Dataset
145
145
-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
146
+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
147
148
```json
149
{"prompt": "What is the capital of India?"}
@@ -173,9 +174,62 @@ vllm bench serve --port 9001 --save-result --save-detailed \
174
175
You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`.
176
176
-#### Custom multimodal dataset
177
+#### Custom Audio Dataset
178
178
-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`:
179
+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`:
180
+
181
+```json
182
+{"prompt": "What does this audio say?", "audio": "/path/to/audio_1.wav"}
183
+{"prompt": "Transcribe the audio.", "audio": "/path/to/audio_2.wav"}
184
+```
185
+
186
+- **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.
187
+
188
+- **Example 1: Whisper**
189
+
190
+Whisper is a dedicated ASR encoder-decoder model, so it uses `--backend openai-audio` and `--endpoint /v1/audio/transcriptions`.
191
+
192
+```bash
193
+# start server
194
+vllm serve openai/whisper-tiny
195
+```
196
+
197
+```bash
198
+vllm bench serve \
199
+ --model openai/whisper-tiny \
200
+ --backend openai-audio \
201
+ --endpoint /v1/audio/transcriptions \
202
+ --dataset-name custom_audio \
203
+ --dataset-path audio_data.jsonl \
204
+ --no-oversample \
205
+ --custom-output-len 256 \
206
+ --save-result \
207
+ --save-detailed \
208
+ --result-filename whisper_bench.json
209
+```
210
+
211
+- **Example 2: Qwen2-Audio**
212
+
213
+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.
214
+
215
+```bash
216
+vllm bench serve \
217
+ --model Qwen/Qwen2-Audio-7B-Instruct \
218
+ --backend openai-chat \
219
+ --endpoint /v1/chat/completions \
220
+ --dataset-name custom_audio \
221
+ --dataset-path audio_data.jsonl \
222
+ --no-oversample \
223
+ --custom-output-len 256 \
224
+ --enable-multimodal-chat \
225
+ --save-result \
226
+ --save-detailed \
227
+ --result-filename qwen_bench.json
228
+```
229
+
230
+#### Custom Image Dataset
231
+
232
+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`:
233
234
```json
235
{"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 \
247
--backend openai-chat \
248
--model Qwen/Qwen2-VL-7B-Instruct \
249
--endpoint /v1/chat/completions \
196
- --dataset-name custom_mm \
197
- --dataset-path <path-to-your-mm-data-jsonl> \
250
+ --dataset-name custom_image \
251
+ --dataset-path <path-to-your-image-data-jsonl> \
252
--allowed-local-media-path /path/to/image/folder
253
```
254
vllm/benchmarks/datasets/__init__.py
+6
-2
@@ -9,8 +9,9 @@ from vllm.benchmarks.datasets.datasets import (
9
BlazeditDataset,
10
BurstGPTDataset,
11
ConversationDataset,
12
+ CustomAudioDataset,
13
CustomDataset,
13
- CustomMMDataset,
14
+ CustomImageDataset,
15
HuggingFaceDataset,
16
InstructCoderDataset,
17
MLPerfDataset,
@@ -36,6 +37,7 @@ from vllm.benchmarks.datasets.datasets import (
37
is_valid_sequence,
38
lora_path_on_disk,
39
lora_tokenizer_cache,
40
+ process_audio,
41
process_image,
42
process_video,
43
zeta_prompt,
@@ -51,7 +53,8 @@ __all__ = [
53
"BurstGPTDataset",
54
"ConversationDataset",
55
"CustomDataset",
54
- "CustomMMDataset",
56
+ "CustomAudioDataset",
57
+ "CustomImageDataset",
58
"HuggingFaceDataset",
59
"InstructCoderDataset",
60
"MLPerfDataset",
@@ -77,6 +80,7 @@ __all__ = [
80
"is_valid_sequence",
81
"lora_path_on_disk",
82
"lora_tokenizer_cache",
83
+ "process_audio",
84
"process_image",
85
"process_video",
86
"RangeRatio",
vllm/benchmarks/datasets/datasets.py
+150
-4
@@ -60,6 +60,11 @@ try:
60
except ImportError:
61
pd = PlaceholderModule("pandas")
62
63
+try:
64
+ import soundfile as sf
65
+except ImportError:
66
+ sf = PlaceholderModule("soundfile")
67
+
68
69
logger = logging.getLogger(__name__)
70
@@ -441,6 +446,27 @@ def process_video(video: Any) -> Mapping[str, Any]:
446
)
447
448
449
+def process_audio(audio: Any) -> tuple:
450
+ """
451
+ Process a single audio input and return a (array, sample_rate) tuple.
452
+
453
+ Supports:
454
+ 1. String: treated as a file path, loaded with soundfile.
455
+ 2. Dict with 'array' and 'sampling_rate' keys: HuggingFace audio format.
456
+ 3. Tuple (array, sr): passed through directly.
457
+ """
458
+ if isinstance(audio, str):
459
+ return sf.read(audio)
460
+ if isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
461
+ return audio["array"], audio["sampling_rate"]
462
+ if isinstance(audio, tuple) and len(audio) == 2:
463
+ return audio
464
+ raise ValueError(
465
+ f"Invalid audio input {audio}. Must be a file path string, "
466
+ "a dict with 'array' and 'sampling_rate', or a (array, sr) tuple."
467
+ )
468
+
469
+
470
def gen_prompt_decode_to_target_len(
471
tokenizer: TokenizerLike,
472
token_sequence: list[int],
@@ -1399,6 +1425,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1425
"random-rerank",
1426
"hf",
1427
"custom",
1428
+ "custom_audio",
1429
+ "custom_image",
1430
"custom_mm",
1431
"prefix_repetition",
1432
"spec_bench",
@@ -1816,8 +1844,28 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
1844
no_oversample=args.no_oversample,
1845
)
1846
1819
- elif args.dataset_name == "custom_mm":
1820
- dataset = CustomMMDataset(
1847
+ elif args.dataset_name in ("custom_image", "custom_mm"):
1848
+ if args.dataset_name == "custom_mm":
1849
+ logger.warning(
1850
+ "Dataset name 'custom_mm' is deprecated and will be removed in v0.24. "
1851
+ "Use '--dataset-name custom_image' instead."
1852
+ )
1853
+ dataset = CustomImageDataset(
1854
+ dataset_path=args.dataset_path,
1855
+ disable_shuffle=args.disable_shuffle,
1856
+ random_seed=args.seed,
1857
+ )
1858
+ input_requests = dataset.sample(
1859
+ num_requests=args.num_prompts,
1860
+ tokenizer=tokenizer,
1861
+ output_len=args.custom_output_len,
1862
+ enable_multimodal_chat=args.enable_multimodal_chat,
1863
+ request_id_prefix=args.request_id_prefix,
1864
+ no_oversample=args.no_oversample,
1865
+ )
1866
+
1867
+ elif args.dataset_name == "custom_audio":
1868
+ dataset = CustomAudioDataset(
1869
dataset_path=args.dataset_path,
1870
disable_shuffle=args.disable_shuffle,
1871
random_seed=args.seed,
@@ -2249,9 +2297,9 @@ class CustomDataset(BenchmarkDataset):
2297
return sampled_requests
2298
2299
2252
-class CustomMMDataset(CustomDataset):
2300
+class CustomImageDataset(CustomDataset):
2301
"""
2254
- Implements the Custom MultiModal dataset. Loads data from a JSONL file and generates
2302
+ Implements the Custom image dataset. Loads data from a JSONL file and generates
2303
sample requests based on conversation turns. E.g.,
2304
```
2305
{
@@ -2328,6 +2376,104 @@ class CustomMMDataset(CustomDataset):
2376
return sampled_requests
2377
2378
2379
+class CustomAudioDataset(CustomDataset):
2380
+ """
2381
+ Custom dataset for audio benchmarking. Loads data from a JSONL file. E.g.,
2382
+ {"prompt": "Transcribe the audio.", "audio": "/path/to/audio.wav"}
2383
+
2384
+ Supports both:
2385
+ - Dedicated ASR models (e.g. Whisper) via openai-audio & /v1/audio/transcriptions
2386
+ - Chat-based audio models (e.g. Qwen2-Audio) via openai-chat & /v1/chat/completions
2387
+ """
2388
+
2389
+ IS_MULTIMODAL = True
2390
+
2391
+ def sample(
2392
+ self,
2393
+ tokenizer: TokenizerLike,
2394
+ num_requests: int,
2395
+ output_len: int | None = None,
2396
+ request_id_prefix: str = "",
2397
+ no_oversample: bool = False,
2398
+ skip_chat_template: bool = False,
2399
+ enable_multimodal_chat: bool = False,
2400
+ **kwargs,
2401
+ ) -> list[SampleRequest]:
2402
+ self.num_available_samples = len(self.data)
2403
+ if num_requests <= 0:
2404
+ num_requests = self.num_available_samples
2405
+ sampled_requests = []
2406
+ for i, item in enumerate(self.data):
2407
+ if len(sampled_requests) >= num_requests:
2408
+ break
2409
+ prompt = item.get("prompt", "")
2410
+ if tokenizer is None:
2411
+ prompt_len = 1
2412
+ new_output_len = output_len if output_len not in (None, -1) else 256
2413
+ mm_content = None
2414
+ else:
2415
+ use_chat_template = (
2416
+ not skip_chat_template
2417
+ and hasattr(tokenizer, "chat_template")
2418
+ and tokenizer.chat_template is not None
2419
+ )
2420
+ if enable_multimodal_chat:
2421
+ # Chat-based audio models (e.g., Qwen2-Audio):
2422
+ # encode audio as base64; serve.py assembles the chat message
2423
+ # as: {"role": "user", "content": [
2424
+ # {"type": "text", "text": prompt},
2425
+ # {"type": "input_audio", "input_audio": {...}}
2426
+ # ]}
2427
+ y, sr = process_audio(item["audio"])
2428
+ buf = io.BytesIO()
2429
+ sf.write(buf, y, sr, format="WAV")
2430
+ audio_base64 = base64.b64encode(buf.getvalue()).decode("utf-8")
2431
+ mm_content = {
2432
+ "type": "input_audio",
2433
+ "input_audio": {
2434
+ "data": audio_base64,
2435
+ "format": "wav",
2436
+ },
2437
+ }
2438
+ # prompt stays as plain string; serve.py handles wrapping
2439
+ else:
2440
+ # Whisper-style models: load audio array locally
2441
+ y, sr = process_audio(item["audio"])
2442
+ mm_content = {"audio": (y, sr)}
2443
+ if use_chat_template:
2444
+ # ASR models with a chat template but not multimodal chat
2445
+ prompt = tokenizer.apply_chat_template(
2446
+ [{"role": "user", "content": prompt}],
2447
+ add_generation_prompt=True,
2448
+ tokenize=False,
2449
+ )
2450
+ # else: plain prompt for Whisper-style models
2451
+ prompt_len = (
2452
+ len(tokenizer(prompt).input_ids) if isinstance(prompt, str) else 1
2453
+ )
2454
+ new_output_len = output_len
2455
+ if output_len is None or output_len == -1:
2456
+ if "output_tokens" not in item:
2457
+ raise ValueError(
2458
+ "If no output length is provided the "
2459
+ "custom dataset must contain an 'output_tokens' field."
2460
+ )
2461
+ new_output_len = int(item["output_tokens"])
2462
+ sampled_requests.append(
2463
+ SampleRequest(
2464
+ prompt=prompt,
2465
+ prompt_len=prompt_len,
2466
+ expected_output_len=new_output_len,
2467
+ multi_modal_data=mm_content,
2468
+ request_id=request_id_prefix + str(i),
2469
+ )
2470
+ )
2471
+ self.maybe_oversample_requests(
2472
+ sampled_requests, num_requests, request_id_prefix, no_oversample
2473
+ )
2474
+ return sampled_requests
2475
+
2476
+
2477
# -----------------------------------------------------------------------------
2478
# Spec Bench Dataset Implementation
2479
# -----------------------------------------------------------------------------