From 28ee78af543c563a2fbf78829a7688120e4e4eb5 Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Tue, 12 May 2026 05:17:58 +0100 Subject: [PATCH] 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> --- docs/benchmarking/cli.md | 66 ++++++++++-- vllm/benchmarks/datasets/__init__.py | 8 +- vllm/benchmarks/datasets/datasets.py | 154 ++++++++++++++++++++++++++- 3 files changed, 216 insertions(+), 12 deletions(-) diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 4f0c933c4a..1520e92f3d 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -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 \ + --dataset-name custom_image \ + --dataset-path \ --allowed-local-media-path /path/to/image/folder ``` diff --git a/vllm/benchmarks/datasets/__init__.py b/vllm/benchmarks/datasets/__init__.py index 5d5e172e7b..b989958edc 100644 --- a/vllm/benchmarks/datasets/__init__.py +++ b/vllm/benchmarks/datasets/__init__.py @@ -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", diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index b032c0a0d6..e71b4a3135 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -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 # -----------------------------------------------------------------------------