benchmark_dataset.py 46 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
"""
This module defines a framework for sampling benchmark requests from various
datasets. Each dataset subclass of BenchmarkDataset must implement sample
generation. Supported dataset types include:
  - ShareGPT
  - Random (synthetic)
  - Sonnet
  - BurstGPT
  - HuggingFace
  - VisionArena
"""

import base64
import io
import json
18
import logging
19
20
21
import random
from abc import ABC, abstractmethod
from collections.abc import Mapping
22
from copy import deepcopy
23
24
from dataclasses import dataclass
from functools import cache
25
26
from io import BytesIO
from typing import Any, Callable, Optional, Union
27
28
29
30
31
32
33
34
35
36

import numpy as np
import pandas as pd
from datasets import load_dataset
from PIL import Image
from transformers import PreTrainedTokenizerBase

from vllm.lora.request import LoRARequest
from vllm.lora.utils import get_adapter_absolute_path
from vllm.multimodal import MultiModalDataDict
37
from vllm.multimodal.image import convert_image_mode
38
39
from vllm.transformers_utils.tokenizer import AnyTokenizer, get_lora_tokenizer

40
41
logger = logging.getLogger(__name__)

42
43
44
45
46
47
48
49
50
51
52
# -----------------------------------------------------------------------------
# Data Classes
# -----------------------------------------------------------------------------


@dataclass
class SampleRequest:
    """
    Represents a single inference request for benchmarking.
    """

53
    prompt: Union[str, Any]
54
55
    prompt_len: int
    expected_output_len: int
56
    multi_modal_data: Optional[Union[MultiModalDataDict, dict, list[dict]]] = None
57
    lora_request: Optional[LoRARequest] = None
58
    request_id: Optional[str] = None
59
60
61
62
63
64
65
66
67


# -----------------------------------------------------------------------------
# Benchmark Dataset Base Class
# -----------------------------------------------------------------------------


class BenchmarkDataset(ABC):
    DEFAULT_SEED = 0
68
    IS_MULTIMODAL = False
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

    def __init__(
        self,
        dataset_path: Optional[str] = None,
        random_seed: int = DEFAULT_SEED,
    ) -> None:
        """
        Initialize the BenchmarkDataset with an optional dataset path and random
        seed.  Args:
            dataset_path (Optional[str]): Path to the dataset. If None, it
            indicates that a default or random dataset might be used.
            random_seed (int): Seed value for reproducible shuffling or
            sampling. Defaults to DEFAULT_SEED.
        """
        self.dataset_path = dataset_path
        # Set the random seed, ensuring that a None value is replaced with the
        # default seed.
86
        self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED
87
88
        self.data = None

89
    def apply_multimodal_chat_transformation(
90
91
        self, prompt: str, mm_content: Optional[MultiModalDataDict] = None
    ) -> list[dict]:
92
93
        """
        Transform a prompt and optional multimodal content into a chat format.
94
95
        This method is used for chat models that expect a specific conversation
        format.
96
97
98
99
100
101
        """
        content = [{"text": prompt, "type": "text"}]
        if mm_content is not None:
            content.append(mm_content)
        return [{"role": "user", "content": content}]

102
103
104
    def load_data(self) -> None:
        """
        Load data from the dataset path into self.data.
105

106
107
        This method must be overridden by subclasses since the method to load
        data will vary depending on the dataset format and source.
108

109
110
111
112
        Raises:
            NotImplementedError: If a subclass does not implement this method.
        """
        # TODO (jenniferzhao): add support for downloading data
113
        raise NotImplementedError("load_data must be implemented in subclasses.")
114
115
116
117
118
119
120
121
122
123

    def get_random_lora_request(
        self,
        tokenizer: PreTrainedTokenizerBase,
        max_loras: Optional[int] = None,
        lora_path: Optional[str] = None,
    ) -> tuple[Optional[LoRARequest], AnyTokenizer]:
        """
        Optionally select a random LoRA request and return its associated
        tokenizer.
124

125
126
127
        This method is used when LoRA parameters are provided.  It randomly
        selects a LoRA based on max_loras and retrieves a cached tokenizer for
        that LoRA if available. Otherwise, it returns the base tokenizer.
128

129
130
131
132
133
134
        Args:
            tokenizer (PreTrainedTokenizerBase): The base tokenizer to use if no
            LoRA is selected.  max_loras (Optional[int]): The maximum number of
            LoRAs available. If None, LoRA is not used.  lora_path
            (Optional[str]): Path to the LoRA parameters on disk. If None, LoRA
            is not used.
135

136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
        Returns:
            tuple[Optional[LoRARequest], AnyTokenizer]: A tuple where the first
            element is a LoRARequest (or None if not applicable) and the second
            element is the tokenizer associated with the LoRA request (or the
            base tokenizer).
        """
        if max_loras is None or lora_path is None:
            return None, tokenizer

        # Generate a random LoRA ID in the range [1, max_loras].
        lora_id = random.randint(1, max_loras)
        lora_request = LoRARequest(
            lora_name=str(lora_id),
            lora_int_id=lora_id,
            lora_path=lora_path_on_disk(lora_path),
        )
        if lora_id not in lora_tokenizer_cache:
            lora_tokenizer_cache[lora_id] = get_lora_tokenizer(lora_request)
        # Return lora_request and the cached tokenizer if available; otherwise,
        # return the base tokenizer
        return lora_request, lora_tokenizer_cache[lora_id] or tokenizer

    @abstractmethod
159
    def sample(
160
161
162
163
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        request_id_prefix: str = "",
164
    ) -> list[SampleRequest]:
165
166
        """
        Abstract method to generate sample requests from the dataset.
167

168
169
        Subclasses must override this method to implement dataset-specific logic
        for generating a list of SampleRequest objects.
170

171
172
173
174
        Args:
            tokenizer (PreTrainedTokenizerBase): The tokenizer to be used
             for processing the dataset's text.
            num_requests (int): The number of sample requests to generate.
175
            request_id_prefix (str) The prefix of request_id.
176

177
178
179
180
181
182
        Returns:
            list[SampleRequest]: A list of sample requests generated from the
            dataset.
        """
        raise NotImplementedError("sample must be implemented in subclasses.")

183
    def maybe_oversample_requests(
184
185
186
187
        self,
        requests: list[SampleRequest],
        num_requests: int,
        request_id_prefix: str = "",
188
    ) -> None:
189
190
191
192
193
194
        """
        Oversamples the list of requests if its size is less than the desired
        number.

        Args:
            requests (List[SampleRequest]): The current list of sampled
195
196
197
            requests.
            num_requests (int): The target number of requests.
            request_id_prefix (str) The prefix of the request ids.
198
199
200
        """
        if len(requests) < num_requests:
            random.seed(self.random_seed)
201
202
203
204
205
206
            additional = deepcopy(
                random.choices(requests, k=num_requests - len(requests))
            )
            for i in range(len(additional)):
                req = additional[i]
                req.request_id = request_id_prefix + str(len(requests) + i)
207
            requests.extend(additional)
208
            logger.info("Oversampled requests to reach %d total samples.", num_requests)
209

210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232

# -----------------------------------------------------------------------------
# Utility Functions and Global Caches
# -----------------------------------------------------------------------------


def is_valid_sequence(
    prompt_len: int,
    output_len: int,
    min_len: int = 4,
    max_prompt_len: int = 1024,
    max_total_len: int = 2048,
    skip_min_output_len_check: bool = False,
) -> bool:
    """
    Validate a sequence based on prompt and output lengths.

    Default pruning criteria are copied from the original `sample_hf_requests`
    and `sample_sharegpt_requests` functions in benchmark_serving.py, as well as
    from `sample_requests` in benchmark_throughput.py.
    """
    # Check for invalid conditions
    prompt_too_short = prompt_len < min_len
233
    output_too_short = (not skip_min_output_len_check) and (output_len < min_len)
234
235
236
237
    prompt_too_long = prompt_len > max_prompt_len
    combined_too_long = (prompt_len + output_len) > max_total_len

    # Return True if none of the invalid conditions are met
238
239
240
    return not (
        prompt_too_short or output_too_short or prompt_too_long or combined_too_long
    )
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255


@cache
def lora_path_on_disk(lora_path: str) -> str:
    return get_adapter_absolute_path(lora_path)


# Global cache for LoRA tokenizers.
lora_tokenizer_cache: dict[int, AnyTokenizer] = {}


def process_image(image: Any) -> Mapping[str, Any]:
    """
    Process a single image input and return a multimedia content dictionary.

256
    Supports three input types:
257

258
259
260
261
262
263
264
265
266
267
    1. Dictionary with raw image bytes: - Expects a dict with a 'bytes' key
       containing raw image data.  - Loads the bytes as a PIL.Image.Image.

    2. PIL.Image.Image input: - Converts the image to RGB.  - Saves the image as
       a JPEG in memory.  - Encodes the JPEG data as a base64 string.  - Returns
       a dictionary with the image as a base64 data URL.

    3. String input: - Treats the string as a URL or local file path.  -
       Prepends "file://" if the string doesn't start with "http://" or
       "file://".  - Returns a dictionary with the image URL.
268
269

    Raises:
270
        ValueError: If the input is not a supported type.
271
    """
272
273
    if isinstance(image, dict) and "bytes" in image:
        image = Image.open(BytesIO(image["bytes"]))
274
    if isinstance(image, Image.Image):
275
        image = convert_image_mode(image, "RGB")
276
277
        with io.BytesIO() as image_data:
            image.save(image_data, format="JPEG")
278
            image_base64 = base64.b64encode(image_data.getvalue()).decode("utf-8")
279
280
        return {
            "type": "image_url",
281
            "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"},
282
283
284
        }

    if isinstance(image, str):
285
286
287
        image_url = (
            image if image.startswith(("http://", "file://")) else f"file://{image}"
        )
288
289
        return {"type": "image_url", "image_url": {"url": image_url}}

290
291
292
293
    raise ValueError(
        f"Invalid image input {image}. Must be a PIL.Image.Image"
        " or str or dictionary with raw image bytes."
    )
294
295


296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
def process_video(video: Any) -> Mapping[str, Any]:
    """
    Process a single video input and return a multimedia content dictionary.

    Supports the following input types:

    1. Dictionary with raw video bytes: - Expects a dict with a 'bytes' key
       containing raw video data.

    2. String input: - Treats the string as a URL or local file path.  -
       Prepends "file://" if the string doesn't start with "http://" or
       "file://".  - Returns a dictionary with the image URL.

    Raises:
        ValueError: If the input is not a supported type.
    """
    if isinstance(video, dict) and "bytes" in video:
        video_bytes = video["bytes"]
        video_base64 = base64.b64encode(video_bytes).decode("utf-8")
        return {
            "type": "video_url",
            "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
        }

    if isinstance(video, str):
        video_url = (
            video if video.startswith(("http://", "file://")) else f"file://{video}"
        )
        return {"type": "video_url", "video_url": {"url": video_url}}

    raise ValueError(
        f"Invalid video input {video}. Must be a string of local path/remote url, or a dictionary with raw video bytes in the form of `{{'bytes': raw_video_bytes}}`."  # noqa: E501
    )


331
332
333
334
335
336
337
338
# -----------------------------------------------------------------------------
# Random Dataset Implementation (Synthetic Data)
# -----------------------------------------------------------------------------


class RandomDataset(BenchmarkDataset):
    # Default values copied from benchmark_serving.py for the random dataset.
    DEFAULT_PREFIX_LEN = 0
339
    DEFAULT_RANGE_RATIO = 0.0
340
341
342
343
344
345
346
347
348
    DEFAULT_INPUT_LEN = 1024
    DEFAULT_OUTPUT_LEN = 128

    def __init__(
        self,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

349
350
351
352
353
354
355
356
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        range_ratio: float = DEFAULT_RANGE_RATIO,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
357
        request_id_prefix: str = "",
358
359
        **kwargs,
    ) -> list[SampleRequest]:
360
361
362
363
364
        # Enforce range_ratio < 1
        assert range_ratio < 1.0, (
            "random_range_ratio must be < 1.0 to ensure a valid sampling range"
        )

365
        vocab_size = tokenizer.vocab_size
366
367
        num_special_tokens = tokenizer.num_special_tokens_to_add()
        real_input_len = input_len - num_special_tokens
368

369
370
371
372
373
        prefix_token_ids = (
            np.random.randint(0, vocab_size, size=prefix_len).tolist()
            if prefix_len > 0
            else []
        )
374

375
        # New sampling logic: [X * (1 - b), X * (1 + b)]
376
377
        input_low = int(real_input_len * (1 - range_ratio))
        input_high = int(real_input_len * (1 + range_ratio))
378
        output_low = int(output_len * (1 - range_ratio))
379
380
381
        # Ensure the lower bound for output length is at least 1 to prevent
        # sampling 0 tokens, which can cause request failures.
        output_low = max(output_low, 1)
382
383
384
385
        output_high = int(output_len * (1 + range_ratio))

        # Add logging for debugging
        logger.info("Sampling input_len from [%s, %s]", input_low, input_high)
386
387
388
389
        logger.info("Sampling output_len from [%s, %s]", output_low, output_high)

        input_lens = np.random.randint(input_low, input_high + 1, size=num_requests)
        output_lens = np.random.randint(output_low, output_high + 1, size=num_requests)
390
391
392
393
        offsets = np.random.randint(0, vocab_size, size=num_requests)

        requests = []
        for i in range(num_requests):
394
395
396
            inner_seq = (
                (offsets[i] + i + np.arange(input_lens[i])) % vocab_size
            ).tolist()
397
398
            token_sequence = prefix_token_ids + inner_seq
            prompt = tokenizer.decode(token_sequence)
399
400
401
402
403
404
405
            # After decoding the prompt we have to encode and decode it again.
            # This is done because in some cases N consecutive tokens
            # give a string tokenized into != N number of tokens.
            # For example for GPT2Tokenizer:
            # [6880, 6881] -> ['Ġcalls', 'here'] ->
            # [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']
            # To avoid uncontrolled change of the prompt length,
406
            # the encoded sequence is truncated before being decoded again.
407
            total_input_len = prefix_len + int(input_lens[i])
408
            re_encoded_sequence = tokenizer.encode(prompt, add_special_tokens=False)[
409
                :total_input_len
410
            ]
411
            prompt = tokenizer.decode(re_encoded_sequence)
412
            total_input_len = len(re_encoded_sequence)
413
414
415
416
417
            requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=total_input_len,
                    expected_output_len=int(output_lens[i]),
418
                    request_id=request_id_prefix + str(i),
419
420
                )
            )
421

422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
        return requests


# -----------------------------------------------------------------------------
# ShareGPT Dataset Implementation
# -----------------------------------------------------------------------------


class ShareGPTDataset(BenchmarkDataset):
    """
    Implements the ShareGPT dataset.  Loads data from a JSON file and generates
    sample requests based on conversation turns.
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.load_data()

    def load_data(self) -> None:
        if self.dataset_path is None:
            raise ValueError("dataset_path must be provided for loading data.")

        with open(self.dataset_path, encoding="utf-8") as f:
            self.data = json.load(f)
        # Filter entries with at least two conversation turns.
        self.data = [
448
449
            entry
            for entry in self.data
450
451
452
453
454
            if "conversations" in entry and len(entry["conversations"]) >= 2
        ]
        random.seed(self.random_seed)
        random.shuffle(self.data)

455
456
457
458
459
460
461
462
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        lora_path: Optional[str] = None,
        max_loras: Optional[int] = None,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
463
        request_id_prefix: str = "",
464
465
        **kwargs,
    ) -> list:
466
        samples: list = []
467
        ind = 0
468
469
470
        for entry in self.data:
            if len(samples) >= num_requests:
                break
471
472
473
474
            prompt, completion = (
                entry["conversations"][0]["value"],
                entry["conversations"][1]["value"],
            )
475
476

            lora_request, tokenizer = self.get_random_lora_request(
477
478
                tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path
            )
479
480
481
            prompt_ids = tokenizer(prompt).input_ids
            completion_ids = tokenizer(completion).input_ids
            prompt_len = len(prompt_ids)
482
483
484
485
486
487
            new_output_len = len(completion_ids) if output_len is None else output_len
            if not is_valid_sequence(
                prompt_len,
                new_output_len,
                skip_min_output_len_check=output_len is not None,
            ):
488
                continue
489
490
            if image_path := entry.get("image"):
                mm_content = process_image(image_path)
491
492
            elif video_path := entry.get("video"):
                mm_content = process_video(video_path)
493
494
            else:
                mm_content = None
495
            if enable_multimodal_chat:
496
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
497
498
499
500
501
502
            samples.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=new_output_len,
                    lora_request=lora_request,
503
                    multi_modal_data=mm_content,
504
                    request_id=request_id_prefix + str(ind),
505
506
                )
            )
507
508
            ind += 1
        self.maybe_oversample_requests(samples, num_requests, request_id_prefix)
509
510
511
        return samples


512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
# -----------------------------------------------------------------------------
# Custom Dataset Implementation
# -----------------------------------------------------------------------------


class CustomDataset(BenchmarkDataset):
    """
    Implements the Custom dataset.  Loads data from a JSONL file and generates
    sample requests based on conversation turns. E.g.,
    ```
    {"prompt": "What is the capital of India?"}
    {"prompt": "What is the capital of Iran?"}
    {"prompt": "What is the capital of China?"}
    ```
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.load_data()

    def load_data(self) -> None:
        if self.dataset_path is None:
            raise ValueError("dataset_path must be provided for loading data.")

        # self.data will be a list of dictionaries
        # e.g., [{"prompt": "What is the capital of India?"}, ...]
        # This will be the standardized format which load_data()
        # has to convert into depending on the filetype of dataset_path.
        # sample() will assume this standardized format of self.data
        self.data = []

        # Load the JSONL file
        if self.dataset_path.endswith(".jsonl"):
            jsonl_data = pd.read_json(path_or_buf=self.dataset_path, lines=True)

            # check if the JSONL file has a 'prompt' column
            if "prompt" not in jsonl_data.columns:
                raise ValueError("JSONL file must contain a 'prompt' column.")

            # Convert each row to a dictionary and append to self.data
            # This will convert the DataFrame to a list of dictionaries
            # where each dictionary corresponds to a row in the DataFrame.
            # This is the standardized format we want for self.data
            for _, row in jsonl_data.iterrows():
                self.data.append(row.to_dict())
        else:
            raise NotImplementedError(
                "Only JSONL format is supported for CustomDataset."
            )

        random.seed(self.random_seed)
        random.shuffle(self.data)

    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        lora_path: Optional[str] = None,
        max_loras: Optional[int] = None,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
574
        request_id_prefix: str = "",
575
576
577
        **kwargs,
    ) -> list:
        sampled_requests = []
578
        for i, item in enumerate(self.data):
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
            if len(sampled_requests) >= num_requests:
                break
            prompt = item["prompt"]

            # apply template
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
                    [{"role": "user", "content": prompt}],
                    add_generation_prompt=True,
                    tokenize=False,
                )

            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
597
                    request_id=request_id_prefix + str(i),
598
599
                )
            )
600
601
602
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
603
604
605
606

        return sampled_requests


607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
# -----------------------------------------------------------------------------
# Sonnet Dataset Implementation
# -----------------------------------------------------------------------------


class SonnetDataset(BenchmarkDataset):
    """
    Simplified implementation of the Sonnet dataset.  Loads poem lines from a
    text file and generates sample requests.  Default values here copied from
    `benchmark_serving.py` for the sonnet dataset.
    """

    DEFAULT_PREFIX_LEN = 200
    DEFAULT_INPUT_LEN = 550
    DEFAULT_OUTPUT_LEN = 150

    def __init__(
        self,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.load_data()

    def load_data(self) -> None:
        if not self.dataset_path:
            raise ValueError("dataset_path must be provided.")
        with open(self.dataset_path, encoding="utf-8") as f:
            self.data = f.readlines()

636
637
638
639
640
641
642
643
    def sample(
        self,
        tokenizer,
        num_requests: int,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        return_prompt_formatted: bool = False,
644
        request_id_prefix: str = "",
645
646
        **kwargs,
    ) -> list:
647
648
        # Calculate average token length for a poem line.
        tokenized_lines = [tokenizer(line).input_ids for line in self.data]
649
        avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)
650
651
652
653

        # Build the base prompt.
        base_prompt = "Pick as many lines as you can from these poem lines:\n"
        base_msg = [{"role": "user", "content": base_prompt}]
654
655
656
        base_fmt = tokenizer.apply_chat_template(
            base_msg, add_generation_prompt=True, tokenize=False
        )
657
658
659
660
        base_offset = len(tokenizer(base_fmt).input_ids)
        if input_len <= base_offset:
            raise ValueError(
                f"'input_len' must be higher than the base prompt length "
661
662
                f"({base_offset})."
            )
663
664
665

        # Determine how many poem lines to use.
        num_input_lines = round((input_len - base_offset) / avg_len)
666
        num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0)
667
668
669
        prefix_lines = self.data[:num_prefix_lines]

        samples = []
670
        ind = 0
671
        while len(samples) < num_requests:
672
673
674
            extra_lines = random.choices(
                self.data, k=num_input_lines - num_prefix_lines
            )
675
676
677
            prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"
            msg = [{"role": "user", "content": prompt}]
            prompt_formatted = tokenizer.apply_chat_template(
678
679
                msg, add_generation_prompt=True, tokenize=False
            )
680
            prompt_len = len(tokenizer(prompt_formatted).input_ids)
681

682
683
684
            if prompt_len <= input_len:
                samples.append(
                    SampleRequest(
685
                        prompt=prompt_formatted if return_prompt_formatted else prompt,
686
687
                        prompt_len=prompt_len,
                        expected_output_len=output_len,
688
                        request_id=request_id_prefix + str(ind),
689
690
                    )
                )
691
                ind += 1
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
        return samples


# -----------------------------------------------------------------------------
# BurstGPT Dataset Implementation
# -----------------------------------------------------------------------------


class BurstGPTDataset(BenchmarkDataset):
    """
    Implements the BurstGPT dataset.  Loads data from a CSV file and generates
    sample requests based on synthetic prompt generation. Only rows with Model
    "GPT-4" and positive response tokens are used.
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self.load_data()

711
712
713
    def load_data(
        self,
    ):
714
715
716
717
718
719
720
721
722
723
724
725
726
        if self.dataset_path is None:
            raise ValueError("dataset_path must be provided for loading data.")

        df = pd.read_csv(self.dataset_path)
        # Filter to keep only GPT-4 rows.
        gpt4_df = df[df["Model"] == "GPT-4"]
        # Remove failed requests (where Response tokens is 0 or less).
        gpt4_df = gpt4_df[gpt4_df["Response tokens"] > 0]
        # Sample the desired number of rows.
        self.data = gpt4_df

    def _sample_loaded_data(self, num_requests: int) -> list:
        if num_requests <= len(self.data):
727
            data = self.data.sample(n=num_requests, random_state=self.random_seed)
728
729
730
731
732
733
734
735
736
        else:
            data = self.data.sample(
                n=num_requests,
                random_state=self.random_seed,
                replace=True,
            )
        # Convert the dataframe to a list of lists.
        return data.values.tolist()

737
738
739
740
741
742
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        max_loras: Optional[int] = None,
        lora_path: Optional[str] = None,
743
        request_id_prefix: str = "",
744
745
        **kwargs,
    ) -> list[SampleRequest]:
746
747
748
749
750
751
        samples = []
        data = self._sample_loaded_data(num_requests=num_requests)
        for i in range(num_requests):
            input_len = int(data[i][2])
            output_len = int(data[i][3])
            lora_req, tokenizer = self.get_random_lora_request(
752
753
                tokenizer=tokenizer, max_loras=max_loras, lora_path=lora_path
            )
754
755
756
757
758
759
760
761
762
763
764
            vocab_size = tokenizer.vocab_size
            # Generate a synthetic prompt: a list of token IDs computed as (i +
            # j) modulo vocab_size.
            token_ids = [(i + j) % vocab_size for j in range(input_len)]
            prompt = tokenizer.decode(token_ids)
            samples.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=input_len,
                    expected_output_len=output_len,
                    lora_request=lora_req,
765
                    request_id=request_id_prefix + str(i),
766
767
                )
            )
768
769
770
771
        return samples


# -----------------------------------------------------------------------------
772
# HuggingFace Dataset Base Implementation
773
774
# -----------------------------------------------------------------------------
class HuggingFaceDataset(BenchmarkDataset):
775
776
777
    """Base class for datasets hosted on HuggingFace."""

    SUPPORTED_DATASET_PATHS: Union[set[str], dict[str, Callable]] = set()
778
779
780

    def __init__(
        self,
781
        dataset_path: str,
782
        dataset_split: str,
783
        no_stream: bool = False,
784
785
786
        dataset_subset: Optional[str] = None,
        **kwargs,
    ) -> None:
787
788
        super().__init__(dataset_path=dataset_path, **kwargs)

789
790
        self.dataset_split = dataset_split
        self.dataset_subset = dataset_subset
791
        self.load_stream = not no_stream
792
793
794
        self.load_data()

    def load_data(self) -> None:
795
        """Load data from HuggingFace datasets."""
796
797
798
799
        self.data = load_dataset(
            self.dataset_path,
            name=self.dataset_subset,
            split=self.dataset_split,
800
            streaming=self.load_stream,
801
        )
802
803
804
805
806
807
808
809
810
811
        self.data = self.data.shuffle(seed=self.random_seed)


# -----------------------------------------------------------------------------
# Conversation Dataset Implementation
# -----------------------------------------------------------------------------


class ConversationDataset(HuggingFaceDataset):
    """Dataset for conversation data with multimodal support."""
812

813
    SUPPORTED_DATASET_PATHS = {
814
815
        "lmms-lab/LLaVA-OneVision-Data",
        "Aeala/ShareGPT_Vicuna_unfiltered",
816
    }
817
    IS_MULTIMODAL = True
818

819
820
821
822
823
824
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
825
        request_id_prefix: str = "",
826
827
        **kwargs,
    ) -> list:
828
        # Filter examples with at least 2 conversations
829
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
830
831
        sampled_requests = []
        dynamic_output = output_len is None
832
        ind = 0
833

834
        for item in filtered_data:
835
836
837
838
839
840
841
842
843
844
845
            if len(sampled_requests) >= num_requests:
                break
            conv = item["conversations"]
            prompt, completion = conv[0]["value"], conv[1]["value"]

            prompt_ids = tokenizer(prompt).input_ids
            completion_ids = tokenizer(completion).input_ids
            prompt_len = len(prompt_ids)
            completion_len = len(completion_ids)
            output_len = completion_len if dynamic_output else output_len
            assert isinstance(output_len, int) and output_len > 0
846
            if dynamic_output and not is_valid_sequence(prompt_len, completion_len):
847
                continue
848
            mm_content = process_image(item["image"]) if "image" in item else None
849
850
851
852
            if enable_multimodal_chat:
                # Note: when chat is enabled the request prompt_len is no longer
                # accurate and we will be using request output to count the
                # actual prompt len and output len
853
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
854
855
856
857
858
859
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
860
                    request_id=request_id_prefix + str(ind),
861
862
                )
            )
863
864
865
866
            ind += 1
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
867
868
869
870
871
872
873
874
        return sampled_requests


# -----------------------------------------------------------------------------
# Vision Arena Dataset Implementation
# -----------------------------------------------------------------------------


875
class VisionArenaDataset(HuggingFaceDataset):
876
877
878
879
880
    """
    Vision Arena Dataset.
    """

    DEFAULT_OUTPUT_LEN = 128
881
    SUPPORTED_DATASET_PATHS = {
882
883
        "lmarena-ai/VisionArena-Chat": lambda x: x["conversation"][0][0]["content"],
        "lmarena-ai/vision-arena-bench-v0.1": lambda x: x["turns"][0][0]["content"],
884
    }
885
    IS_MULTIMODAL = True
886

887
888
889
890
891
892
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
893
        request_id_prefix: str = "",
894
895
        **kwargs,
    ) -> list:
896
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
897
        sampled_requests = []
898
        for i, item in enumerate(self.data):
899
900
            if len(sampled_requests) >= num_requests:
                break
901
902
            parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.dataset_path)
            if parser_fn is None:
903
                raise ValueError(f"Unsupported dataset path: {self.dataset_path}")
904
            prompt = parser_fn(item)
905
            mm_content = process_image(item["images"][0])
906
907
908
909
910
            prompt_len = len(tokenizer(prompt).input_ids)
            if enable_multimodal_chat:
                # Note: when chat is enabled the request prompt_len is no longer
                # accurate and we will be using request output to count the
                # actual prompt len
911
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
912
913
914
915
916
917
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
918
                    request_id=request_id_prefix + str(i),
919
920
                )
            )
921
922
923
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
924
        return sampled_requests
925
926
927
928
929
930
931
932
933
934
935
936


# -----------------------------------------------------------------------------
# Instruct Coder Dataset Implementation
# -----------------------------------------------------------------------------


class InstructCoderDataset(HuggingFaceDataset):
    """
    InstructCoder Dataset.
    https://huggingface.co/datasets/likaixin/InstructCoder

937
938
939
    InstructCoder is the dataset designed for general code editing.  It consists
    of 114,239 instruction-input-output triplets, and covers multiple distinct
    code editing scenario.
940
941
942
    """

    DEFAULT_OUTPUT_LEN = 200  # this is the average default output length
943
944
945
    SUPPORTED_DATASET_PATHS = {
        "likaixin/InstructCoder",
    }
946

947
948
949
950
951
952
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
953
        request_id_prefix: str = "",
954
955
956
        **kwargs,
    ) -> list:
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
957
        sampled_requests = []
958
        for i, item in enumerate(self.data):
959
960
            if len(sampled_requests) >= num_requests:
                break
961
962
963
964
            prompt = (
                f"{item['input']}\n\n{item['instruction']} Just output "
                "the code, do not include any explanation."
            )
965
966
967
968
969
970
971

            # apply template
            prompt = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                add_generation_prompt=True,
                tokenize=False,
            )
972
973
974
975
976
977
            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
978
                    request_id=request_id_prefix + str(i),
979
980
                )
            )
981
982
983
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
984
        return sampled_requests
985
986


987
988
989
990
991
992
993
994
995
996
# -----------------------------------------------------------------------------
# MT-Bench Dataset Implementation
# -----------------------------------------------------------------------------


class MTBenchDataset(HuggingFaceDataset):
    """
    MT-Bench Dataset.
    https://huggingface.co/datasets/philschmid/mt-bench

997
    We create a single turn dataset for MT-Bench.
998
999
    This is similar to Spec decoding benchmark setup in vLLM
    https://github.com/vllm-project/vllm/blob/9d98ab5ec/examples/offline_inference/eagle.py#L14-L18
1000
    """  # noqa: E501
1001
1002
1003
1004
1005
1006

    DEFAULT_OUTPUT_LEN = 256  # avg len used in SD bench in vLLM
    SUPPORTED_DATASET_PATHS = {
        "philschmid/mt-bench",
    }

1007
1008
1009
1010
1011
1012
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        output_len: Optional[int] = None,
        enable_multimodal_chat: bool = False,
1013
        request_id_prefix: str = "",
1014
1015
1016
        **kwargs,
    ) -> list:
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
1017
1018
        sampled_requests = []

1019
        for i, item in enumerate(self.data):
1020
1021
            if len(sampled_requests) >= num_requests:
                break
1022
            prompt = item["turns"][0]
1023
1024

            # apply template
1025
1026
1027
1028
1029
            prompt = tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                add_generation_prompt=True,
                tokenize=False,
            )
1030
1031
1032
1033
1034
1035
1036

            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
1037
                    request_id=request_id_prefix + str(i),
1038
1039
                )
            )
1040
1041
1042
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
1043
1044
1045
        return sampled_requests


1046
1047
1048
1049
1050
1051
1052
1053
1054
# -----------------------------------------------------------------------------
# AIMO Dataset Implementation
# -----------------------------------------------------------------------------


class AIMODataset(HuggingFaceDataset):
    """
    Dataset class for processing a AIMO dataset with reasoning questions.
    """
1055

1056
    SUPPORTED_DATASET_PATHS = {
1057
1058
1059
        "AI-MO/aimo-validation-aime",
        "AI-MO/NuminaMath-1.5",
        "AI-MO/NuminaMath-CoT",
1060
1061
    }

1062
1063
1064
1065
1066
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        output_len: Optional[int] = None,
1067
        request_id_prefix: str = "",
1068
1069
        **kwargs,
    ) -> list:
1070
1071
        sampled_requests = []
        dynamic_output = output_len is None
1072
        ind = 0
1073
1074
1075
1076

        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break
1077
            prompt, completion = item["problem"], item["solution"]
1078
1079
1080
1081
1082
1083
1084

            prompt_ids = tokenizer(prompt).input_ids
            completion_ids = tokenizer(completion).input_ids
            prompt_len = len(prompt_ids)
            completion_len = len(completion_ids)
            output_len = completion_len if dynamic_output else output_len
            assert isinstance(output_len, int) and output_len > 0
1085
1086
1087
            if dynamic_output and not is_valid_sequence(
                prompt_len, completion_len, max_prompt_len=2048, max_total_len=32000
            ):
1088
1089
1090
1091
1092
1093
1094
                continue
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=None,
1095
                    request_id=request_id_prefix + str(ind),
1096
1097
                )
            )
1098
1099
1100
1101
            ind += 1
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
1102
        return sampled_requests
1103
1104


1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
# -----------------------------------------------------------------------------
# Next Edit Prediction Dataset Implementation
# -----------------------------------------------------------------------------


zeta_prompt = """### Instruction:
You are a code completion assistant and your task is to analyze user edits and then rewrite an excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking into account the cursor location.

### User Edits:

{}

### User Excerpt:

{}

### Response:

1123
"""  # noqa: E501
1124
1125
1126


def _format_zeta_prompt(
1127
1128
    sample: dict, original_start_marker: str = "<|editable_region_start|>"
) -> dict:
1129
    """Format the zeta prompt for the Next Edit Prediction (NEP) dataset.
1130
1131
1132

    This function formats examples from the NEP dataset
    into prompts and expected outputs. It could be
1133
    further extended to support more NEP datasets.
1134

1135
    Args:
1136
        sample: The dataset sample containing events,
1137
            inputs, and outputs.
1138
1139
        original_start_marker: The marker indicating the
            start of the editable region. Defaults to
1140
            "<|editable_region_start|>".
1141

1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
    Returns:
        A dictionary with the formatted prompts and expected outputs.
    """
    events = sample["events"]
    input = sample["input"]
    output = sample["output"]
    prompt = zeta_prompt.format(events, input)

    # following the original implementation, extract the focused region
    # from the raw output
    output_start_index = output.find(original_start_marker)
    output_focused_region = output[output_start_index:]
    expected_output = output_focused_region

    return {"prompt": prompt, "expected_output": expected_output}


class NextEditPredictionDataset(HuggingFaceDataset):
    """
    Dataset class for processing a Next Edit Prediction dataset.
    """

    SUPPORTED_DATASET_PATHS = {
        "zed-industries/zeta",
    }
    MAPPING_PROMPT_FUNCS = {
        "zed-industries/zeta": _format_zeta_prompt,
    }

1171
1172
1173
1174
1175
1176
1177
    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        request_id_prefix: str = "",
        **kwargs,
    ):
1178
        formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(self.dataset_path)
1179
1180
1181
        if formatting_prompt_func is None:
            raise ValueError(f"Unsupported dataset path: {self.dataset_path}")
        samples = []
1182
        for i, sample in enumerate(self.data):
1183
1184
1185
1186
1187
1188
            sample = formatting_prompt_func(sample)
            samples.append(
                SampleRequest(
                    prompt=sample["prompt"],
                    prompt_len=len(tokenizer(sample["prompt"]).input_ids),
                    expected_output_len=len(
1189
1190
                        tokenizer(sample["expected_output"]).input_ids
                    ),
1191
                    request_id=request_id_prefix + str(i),
1192
1193
                )
            )
1194
1195
            if len(samples) >= num_requests:
                break
1196
        self.maybe_oversample_requests(samples, num_requests, request_id_prefix)
1197
1198
1199
        return samples


1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
# -----------------------------------------------------------------------------
# ASR Dataset Implementation
# -----------------------------------------------------------------------------


class ASRDataset(HuggingFaceDataset):
    """
    Dataset class for processing a ASR dataset for transcription.
    Tested on the following set:

    +----------------+----------------------------------------+--------------------------+-----------------------------+
    | Dataset        | Domain                                 | Speaking Style           | hf-subset                   |
    +----------------+----------------------------------------+--------------------------+-----------------------------+
    | TED-LIUM       | TED talks                              | Oratory                  | release1, release2, release3|
    |                |                                        |                          | release3-speaker-adaptation |
    | VoxPopuli      | European Parliament                    | Oratory                  | en, de, it, fr,  ...        |
    | LibriSpeech    | Audiobook                              | Narrated                 | "LIUM/tedlium"              |
    | GigaSpeech     | Audiobook, podcast, YouTube            | Narrated, spontaneous    | xs, s, m, l, xl, dev, test  |
    | SPGISpeech     | Financial meetings                     | Oratory, spontaneous     | S, M, L, dev, test          |
    | AMI            | Meetings                               | Spontaneous              | ihm, sdm                    |
    +----------------+----------------------------------------+--------------------------+-----------------------------+

1222
1223
    """  # noqa: E501

1224
    SUPPORTED_DATASET_PATHS = {
1225
1226
1227
1228
1229
1230
        "openslr/librispeech_asr",
        "facebook/voxpopuli",
        "LIUM/tedlium",
        "edinburghcstr/ami",
        "speechcolab/gigaspeech",
        "kensho/spgispeech",
1231
1232
1233
1234
1235
1236
    }

    DEFAULT_OUTPUT_LEN = 128
    IS_MULTIMODAL = True

    # TODO Whisper-specific. Abstract interface when more models are supported.
1237
    TRANSCRIPTION_PREAMBLE = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
1238
1239
1240
1241
1242
1243
1244
    skip_long_audios: bool = True

    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        output_len: Optional[int] = None,
1245
        request_id_prefix: str = "",
1246
1247
1248
        **kwargs,
    ) -> list:
        import librosa
1249
1250

        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
1251
1252
1253
1254
        prompt = ASRDataset.TRANSCRIPTION_PREAMBLE
        prompt_len = len(tokenizer(prompt).input_ids)
        sampled_requests = []
        skipped = 0
1255
        ind = 0
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break
            audio = item["audio"]
            y, sr = audio["array"], audio["sampling_rate"]
            duration_s = librosa.get_duration(y=y, sr=sr)
            # Whisper max supported duration
            if self.skip_long_audios and duration_s > 30:
                skipped += 1
                continue

            mm_content = {"audio": (y, sr)}
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
1274
                    request_id=request_id_prefix + str(ind),
1275
1276
                )
            )
1277
            ind += 1
1278
        if skipped:
1279
1280
1281
1282
1283
1284
            logger.warning(
                "%d samples discarded from dataset due to"
                " their length being greater than"
                " what Whisper supports.",
                skipped,
            )
1285
1286
1287
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix
        )
1288
        return sampled_requests