datasets.py 129 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
"""
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
"""
14

15
import argparse
16
import ast
17
18
19
import io
import json
import logging
20
import math
21
22
import random
from abc import ABC, abstractmethod
23
from collections.abc import Callable, Iterator, Mapping
24
from contextlib import suppress
25
from dataclasses import dataclass, replace
26
27
from functools import cache
from io import BytesIO
28
from tempfile import NamedTemporaryFile
29
from typing import Any, cast
30
31

import numpy as np
32
import pybase64 as base64
33
from huggingface_hub import snapshot_download
34
from PIL import Image
35
from typing_extensions import deprecated
36

37
38
39
40
41
from vllm.benchmarks.datasets.utils import (
    RangeRatio,
    _resolve_range_ratios,
    get_sampling_params,
)
42
from vllm.inputs import MultiModalDataDict
43
44
from vllm.lora.request import LoRARequest
from vllm.lora.utils import get_adapter_absolute_path
45
from vllm.multimodal.audio import get_audio_duration
46
from vllm.multimodal.image import convert_image_mode
47
from vllm.tokenizers import TokenizerLike
48
from vllm.utils.argparse_utils import FlexibleArgumentParser
49
from vllm.utils.import_utils import PlaceholderModule
50
51
52
53
54
55
56
57
58
59
60
61

try:
    from datasets import load_dataset
except ImportError:
    datasets = PlaceholderModule("datasets")
    load_dataset = datasets.placeholder_attr("load_dataset")

try:
    import pandas as pd
except ImportError:
    pd = PlaceholderModule("pandas")

62
63
64

logger = logging.getLogger(__name__)

65
66
DEFAULT_NUM_PROMPTS = 1000

67
68
69
70
71
72
73

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

74
    prompt: str | list[str] | list[dict]
75
    prompt_len: int
76
    expected_output_len: int | None
77
78
79
    multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
    lora_request: LoRARequest | None = None
    request_id: str | None = None
80
81
82
83
84
85
86
87
88


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


class BenchmarkDataset(ABC):
    DEFAULT_SEED = 0
89
    IS_MULTIMODAL = False
90
91
92

    def __init__(
        self,
93
        dataset_path: str | None = None,
94
        random_seed: int = DEFAULT_SEED,
95
96
        disable_shuffle: bool = False,
        **kwargs,
97
98
99
    ) -> None:
        """
        Initialize the BenchmarkDataset with an optional dataset path and random
100
101
        seed.

102
103
        Args:
            dataset_path (Optional[str]): Path to the dataset. If None, it
104
                indicates that a default or random dataset might be used.
105
            random_seed (int): Seed value for reproducible shuffling or
106
                sampling. Defaults to DEFAULT_SEED.
107
108
109
110
        """
        self.dataset_path = dataset_path
        # Set the random seed, ensuring that a None value is replaced with the
        # default seed.
111
        self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED
112
        self.disable_shuffle = disable_shuffle
113
        self.data: Any | None = None
114
115

    def apply_multimodal_chat_transformation(
116
117
        self,
        prompt: str,
118
        mm_content: MultiModalDataDict | dict | list[dict] | None = None,
119
    ) -> list[dict]:
120
121
122
123
124
125
126
        """
        Transform a prompt and optional multimodal content into a chat format.
        This method is used for chat models that expect a specific conversation
        format.
        """
        content = [{"text": prompt, "type": "text"}]
        if mm_content is not None:
127
128
129
130
131
            if isinstance(mm_content, list):
                content.extend(cast(list[dict[str, Any]], mm_content))
            elif isinstance(mm_content, dict):
                content.append(mm_content)
            else:
132
                raise TypeError(
133
                    f"Could not process multimodal content of type: {type(mm_content)}"
134
                )
135
136
137
138
139
140
141
142
143
144
145
146
147
        return [{"role": "user", "content": content}]

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

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

        Raises:
            NotImplementedError: If a subclass does not implement this method.
        """
        # TODO (jenniferzhao): add support for downloading data
148
        raise NotImplementedError("load_data must be implemented in subclasses.")
149
150
151

    def get_random_lora_request(
        self,
152
153
154
        max_loras: int | None = None,
        lora_path: str | None = None,
    ) -> LoRARequest | None:
155
        """
156
        Optionally select a random LoRA request.
157
158

        This method is used when LoRA parameters are provided.  It randomly
159
        selects a LoRA based on max_loras.
160
161

        Args:
162
163
164
165
            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.
166
167

        Returns:
168
169
            A new [`LoRARequest`][vllm.lora.request.LoRARequest]
            (or `None` if not applicable).
170
171
        """
        if max_loras is None or lora_path is None:
172
            return None
173
174
175
176
177
178
179
180

        # 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),
        )
181
        return lora_request
182

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
    def get_round_robin_lora_request(
        self,
        index: int,
        max_loras: int | None = None,
        lora_path: str | None = None,
    ) -> LoRARequest | None:
        """
        Optionally select a LoRA request using deterministic round-robin.

        This method cycles through LoRA IDs in order based on the request
        index, providing reproducible LoRA assignment.

        Args:
            index (int): The request index used for round-robin selection.
            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.

        Returns:
            A new [`LoRARequest`][vllm.lora.request.LoRARequest]
            (or `None` if not applicable).
        """
        if max_loras is None or lora_path is None:
            return None

        # Deterministic round-robin: cycle through [1, max_loras]
        lora_id = index % max_loras + 1
        lora_request = LoRARequest(
            lora_name=str(lora_id),
            lora_int_id=lora_id,
            lora_path=lora_path_on_disk(lora_path),
        )
        return lora_request

    def get_lora_request(
        self,
        index: int,
        max_loras: int | None = None,
        lora_path: str | None = None,
        lora_assignment: str = "random",
    ) -> LoRARequest | None:
        """
        Select a LoRA request using the specified assignment strategy.

        Args:
            index (int): The request index (used for round-robin).
            max_loras (Optional[int]): The maximum number of LoRAs available.
            lora_path (Optional[str]): Path to the LoRA parameters on disk.
            lora_assignment (str): Strategy for LoRA selection.
                'random' (default) or 'round-robin'.

        Returns:
            A new [`LoRARequest`][vllm.lora.request.LoRARequest]
            (or `None` if not applicable).
        """
        if lora_assignment == "round-robin":
            return self.get_round_robin_lora_request(
                index=index, max_loras=max_loras, lora_path=lora_path
            )
        return self.get_random_lora_request(max_loras=max_loras, lora_path=lora_path)

245
    @abstractmethod
246
247
    def sample(
        self,
248
        tokenizer: TokenizerLike,
249
250
251
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
252
        **kwargs,
253
    ) -> list[SampleRequest]:
254
255
256
257
258
259
260
        """
        Abstract method to generate sample requests from the dataset.

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

        Args:
261
            tokenizer (TokenizerLike): The tokenizer to be used
262
                for processing the dataset's text.
263
            num_requests (int): The number of sample requests to generate.
264
            request_id_prefix (str): The prefix of request_id.
265
266
267
268
269
270
271

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

272
273
274
275
276
    def maybe_oversample_requests(
        self,
        requests: list[SampleRequest],
        num_requests: int,
        request_id_prefix: str = "",
277
        no_oversample: bool = False,
278
    ) -> None:
279
280
281
282
283
284
        """
        Oversamples the list of requests if its size is less than the desired
        number.

        Args:
            requests (List[SampleRequest]): The current list of sampled
285
286
                requests.
            num_requests (int): The target number of requests.
287
288
            request_id_prefix (str): The prefix applied to generated request
                identifiers.
289

290
        """
291
        if no_oversample:
292
            logger.info("Skipping oversampling. Total samples: %d.", len(requests))
293
294
            return

295
296
        if len(requests) < num_requests:
            random.seed(self.random_seed)
297
298
299
            needed = num_requests - len(requests)
            additional = []
            for i in range(needed):
300
301
302
303
                req = replace(
                    random.choice(requests),
                    request_id=request_id_prefix + str(len(requests) + i),
                )
304
                additional.append(req)
305
            requests.extend(additional)
306
            logger.info("Oversampled requests to reach %d total samples.", num_requests)
307

308
309
        ids = [req.request_id for req in requests]
        if len(ids) != len(set(ids)):
310
311
312
313
314
            raise ValueError(
                "Duplicate request_id found in the sampled "
                "requests. Please ensure that each request_id "
                "is unique."
            )
315

316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338

# -----------------------------------------------------------------------------
# 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
339
    output_too_short = (not skip_min_output_len_check) and (output_len < min_len)
340
341
342
343
    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
344
345
346
    return not (
        prompt_too_short or output_too_short or prompt_too_long or combined_too_long
    )
347
348
349
350
351
352
353
354


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


# Global cache for LoRA tokenizers.
355
lora_tokenizer_cache: dict[int, TokenizerLike] = {}
356
357
358
359
360
361


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

362
    Supports the following input types:
363
364
365
366
367
368
369
370

    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.

371
372
373
374
375
    3. String input: - Treats the string as a URL, local file path, or base64
       encoded data.  - If string starts with "data:image/", treats as base64.
       - If string starts with "http://", "https://", or "file://", treats as URL.
       - Otherwise treats as local file path and prepends "file://".
       - Returns a dictionary with the image URL or base64 data.
376
377
378
379

    Raises:
        ValueError: If the input is not a supported type.
    """
380
381
    if isinstance(image, dict) and "bytes" in image:
        image = Image.open(BytesIO(image["bytes"]))
382
    if isinstance(image, Image.Image):
383
        image = convert_image_mode(image, "RGB")
384
385
        with io.BytesIO() as image_data:
            image.save(image_data, format="JPEG")
386
            image_base64 = base64.b64encode(image_data.getvalue()).decode("utf-8")
387
388
        return {
            "type": "image_url",
389
            "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"},
390
391
392
        }

    if isinstance(image, str):
393
394
        image_url = (
            image
395
            if image.startswith(("http://", "https://", "file://", "data:image/"))
396
397
            else f"file://{image}"
        )
398
399
        return {"type": "image_url", "image_url": {"url": image_url}}

400
    raise ValueError(
401
402
        f"Invalid image input {image}. Must be a PIL.Image.Image, "
        "str (URL, file path, or base64 data URL), or dictionary with raw image bytes."
403
    )
404
405


406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
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.
    """
422
423
    if isinstance(video, dict) and "bytes" in video:
        video_bytes = video["bytes"]
424
425
426
        video_base64 = base64.b64encode(video_bytes).decode("utf-8")
        return {
            "type": "video_url",
427
            "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
428
429
430
        }

    if isinstance(video, str):
431
432
433
434
435
        video_url = (
            video
            if video.startswith(("http://", "https://", "file://"))
            else f"file://{video}"
        )
436
437
438
439
440
441
        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
    )

442
443

def gen_prompt_decode_to_target_len(
444
    tokenizer: TokenizerLike,
445
446
447
448
    token_sequence: list[int],
    target_token_len: int,
    max_retry: int = 10,
    add_special_tokens: bool = False,
449
    rng: np.random.Generator | None = None,
450
) -> tuple[str, list[int], int]:
451
452
453
454
    """
    Ensure decoded-then-encoded prompt length matches the target token length.

    This function decodes an initial token sequence to text and re-encodes it
455
456
    , iteratively adjusting the token sequence length to match a target.
    This is necessary because some tokenizers do not guarantee a 1:1 mapping
457
458
459
460
461
    between consecutive tokens and the decoded-then-encoded sequence length.
    For example, for GPT2Tokenizer:
    [6880, 6881] -> ['Ġcalls', 'here'] ->
    [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']

462
463
464
    Returns a tuple of the final prompt string, the adjusted token sequence,
    and the token mismatch (final_len - target_token_len) if the retry budget
    is exhausted.
465
466
467
468
469
    """
    remain_num_try = max_retry
    token_mismatch = 0
    while True:
        prompt = tokenizer.decode(token_sequence)
470
        token_sequence = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
471
472
473
474
        if remain_num_try <= 0:
            if len(token_sequence) != target_token_len:
                token_mismatch = len(token_sequence) - target_token_len
            break
475

476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
        if len(token_sequence) == target_token_len:
            break
        elif len(token_sequence) < target_token_len:
            if rng is not None:
                extra_tokens = rng.integers(
                    0,
                    tokenizer.vocab_size,
                    size=target_token_len - len(token_sequence),
                ).tolist()
            else:
                extra_tokens = np.random.randint(
                    0,
                    tokenizer.vocab_size,
                    size=target_token_len - len(token_sequence),
                ).tolist()
            token_sequence.extend(extra_tokens)
        elif len(token_sequence) > target_token_len:
            token_sequence = token_sequence[:target_token_len]

        remain_num_try -= 1

    return prompt, token_sequence, token_mismatch

499

500
501
502
503
# -----------------------------------------------------------------------------
# Random Dataset Implementation (Synthetic Data)
# -----------------------------------------------------------------------------

504

505
class RandomDataset(BenchmarkDataset):
506
507
508
509
510
511
512
513
514
515
516
517
    """
    Synthetic text-only dataset for serving/throughput benchmarks.

    Strategy:
    - Sample input/output token lengths per request from integer-uniform ranges
      around configured means (controlled by range_ratio).
    - Prepend a fixed random prefix of length prefix_len.
    - Generate the remaining tokens as a reproducible sequence:
      (offset + index + arange(input_len)) % vocab_size.
    - Decode then re-encode/truncate to ensure prompt token counts match.
    - Uses numpy.default_rng seeded with random_seed for reproducible sampling.
    """
518

519
520
521
522
523
524
    # Default values copied from benchmark_serving.py for the random dataset.
    DEFAULT_PREFIX_LEN = 0
    DEFAULT_RANGE_RATIO = 0.0
    DEFAULT_INPUT_LEN = 1024
    DEFAULT_OUTPUT_LEN = 128

525
    def __init__(self, **kwargs) -> None:
526
        super().__init__(**kwargs)
527
528
529
530
        # Use numpy's default_rng for deterministic sampling
        # Do not use random.seed() or np.random.seed() elsewhere in this class.
        # This ensures that the RNG is isolated from global RNG state.
        self._rng = np.random.default_rng(self.random_seed)
531
532
533

    def sample(
        self,
534
        tokenizer: TokenizerLike,
535
        num_requests: int,
536
        request_id_prefix: str = "",
537
        no_oversample: bool = False,
538
        prefix_len: int = DEFAULT_PREFIX_LEN,
539
        range_ratio: RangeRatio = DEFAULT_RANGE_RATIO,
540
541
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
542
        batchsize: int = 1,
543
544
545
        max_loras: int | None = None,
        lora_path: str | None = None,
        lora_assignment: str = "random",
546
547
        **kwargs,
    ) -> list[SampleRequest]:
548
549
        resolved_input_rr, _ = _resolve_range_ratios(range_ratio)

550
551
        num_special = int(tokenizer.num_special_tokens_to_add())
        real_input_len = max(0, int(input_len) - num_special)
552
553
554
        min_sampled_input = math.floor(
            real_input_len * (1.0 - float(resolved_input_rr))
        )
555
556
557
558
        min_total_input = int(prefix_len) + min_sampled_input
        if min_total_input < 1:
            raise ValueError(
                "--random-input-len is too small: with tokenizer special "
559
560
                f"tokens {num_special} and "
                f"input range ratio {resolved_input_rr}, "
561
562
                "the minimum possible total input tokens (prefix + sampled) is "
                f"{min_total_input}. Increase --random-input-len and/or "
563
564
565
566
567
568
569
570
571
572
573
574
                "--random-prefix-len, or decrease the input range ratio "
                "so that prefix_len + floor(max(0, random_input_len - "
                "num_special)) * (1 - input_range_ratio) >= 1."
            )

        input_lens, output_lens, offsets = get_sampling_params(
            self._rng,
            num_requests,
            range_ratio,
            input_len,
            output_len,
            tokenizer,
575
576
577
        )

        vocab_size = tokenizer.vocab_size
578
579
580
581
582
        prohibited_tokens = tokenizer.all_special_ids
        all_tokens = np.arange(vocab_size)
        allowed_tokens = np.array(list(set(all_tokens) - set(prohibited_tokens)))

        # Generate prefix once
583
        prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)
584

585
        requests = []
586
        token_mismatch_total = 0
587
        for i in range(num_requests):
588
            prompt, total_input_len, token_mismatch = self.generate_token_sequence(  # noqa: E501
589
590
591
592
593
594
595
                tokenizer=tokenizer,
                prefix_token_ids=prefix_token_ids,
                prefix_len=prefix_len,
                vocab_size=vocab_size,
                input_len=int(input_lens[i]),
                offset=int(offsets[i]),
                index=i,
596
                allowed_tokens=allowed_tokens,
597
            )
598
            token_mismatch_total += token_mismatch
599
600
601
602
603
604
            lora_req = self.get_lora_request(
                index=i,
                max_loras=max_loras,
                lora_path=lora_path,
                lora_assignment=lora_assignment,
            )
605
606
607
608
609
            requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=total_input_len,
                    expected_output_len=int(output_lens[i]),
610
                    lora_request=lora_req,
611
612
613
                    request_id=request_id_prefix + str(i),
                )
            )
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
        # only used for embeddings benchmark.
        if batchsize > 1:
            batch_requests = []
            # Create batched requests
            for i in range(0, num_requests, batchsize):
                batch = requests[i : i + batchsize]
                batch_requests.append(
                    SampleRequest(
                        prompt=[req.prompt for req in batch],
                        prompt_len=sum(req.prompt_len for req in batch),
                        expected_output_len=0,
                        request_id=request_id_prefix + str(i // batchsize),
                    )
                )
            requests = batch_requests
629

630
631
632
633
634
635
636
637
638
639
640
        if token_mismatch_total != 0:
            sign = "more" if token_mismatch_total > 0 else "fewer"
            logger.warning(
                "Across all generated prompts, there were %d %s tokens "
                "than expected after decoding and re-encoding. This is "
                "expected due to the imperfect nature of the sampling "
                "procedure.",
                abs(token_mismatch_total),
                sign,
            )

641
642
643
        return requests

    def get_prefix(
644
        self,
645
        tokenizer: TokenizerLike,
646
647
        allowed_tokens: np.ndarray,
        prefix_len: int,
648
649
650
651
    ) -> list[int]:
        """
        Get the prefix for the dataset.
        """
652
653
654
655
656
657
658
659
660
661
662
663
        if prefix_len <= 0:
            return []

        prefix_tokens = allowed_tokens[
            self._rng.integers(0, len(allowed_tokens), size=prefix_len)
        ].tolist()
        _, adjusted_tokens, token_mismatch = gen_prompt_decode_to_target_len(
            tokenizer=tokenizer,
            token_sequence=prefix_tokens,
            target_token_len=prefix_len,
            add_special_tokens=False,
            rng=self._rng,
664
        )
665
666
667
668
669
670
671
672
673
674
        if token_mismatch != 0:
            sign = "more" if token_mismatch > 0 else "fewer"
            logger.warning(
                "Prefix tokenization produced %d %s tokens than expected "
                "after decoding and re-encoding. This is expected due to "
                "the imperfect nature of the sampling procedure",
                abs(token_mismatch),
                sign,
            )
        return adjusted_tokens
675

676
677
678
    def generate_token_sequence(
        self,
        *,
679
        tokenizer: TokenizerLike,
680
681
682
683
684
685
        prefix_token_ids: list[int],
        prefix_len: int,
        vocab_size: int,
        input_len: int,
        offset: int,
        index: int,
686
        allowed_tokens: np.ndarray,
687
    ) -> tuple[str, int, int]:
688
689
690
691
692
693
694
695
696
697
        """
        Returns (prompt, total_input_len).

        NOTE: 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,
698
        the encoded sequence is truncated before being decoded again.
699
        """
700
701
702
703
704
        # Build the inner sequence by sampling
        # sequentially from the allowed tokens
        inner_seq = allowed_tokens[
            (offset + index + np.arange(input_len)) % len(allowed_tokens)
        ].tolist()
705
706
707
708
        token_sequence = prefix_token_ids + inner_seq

        # Decode, then re-encode and truncate to preserve token count invariants
        total_input_len = prefix_len + int(input_len)
709
        prompt, adjusted_token_sequence, token_mismatch = (
710
            gen_prompt_decode_to_target_len(
711
712
713
714
715
716
                tokenizer=tokenizer,
                token_sequence=token_sequence,
                target_token_len=total_input_len,
                add_special_tokens=False,
                rng=self._rng,
            )
717
718
719
        )
        total_input_len = len(adjusted_token_sequence)
        return prompt, total_input_len, token_mismatch
720
721


722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
# -----------------------------------------------------------------------------
# Random Dataset Implementation (Synthetic Data)
# -----------------------------------------------------------------------------


class RandomDatasetForReranking(RandomDataset):
    """
    Random dataset specialized for the needs of scoring:
    - Batches of inputs
    - Inputs composed of pairs
    """

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

    def sample(
        self,
739
        tokenizer: TokenizerLike,
740
741
        num_requests: int,
        request_id_prefix: str = "",
742
743
744
        no_oversample: bool = False,
        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
745
        input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
746
        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
747
748
749
750
751
752
753
754
        batchsize: int = 1,
        is_reranker: bool = True,
        **kwargs,
    ) -> list[SampleRequest]:
        n_sep_tokens = int(is_reranker)

        query_len_param = (input_len // 2) - n_sep_tokens if is_reranker else input_len

755
756
757
758
759
760
761
        query_lens, _, query_offsets = get_sampling_params(
            self._rng,
            1,
            range_ratio,
            query_len_param,
            0,
            tokenizer,
762
763
764
765
766
767
768
769
770
771
772
773
        )

        query_len = int(query_lens[0])

        if not is_reranker:
            assert num_requests > 1 and batchsize > 1
            num_requests -= 1
            batchsize -= 1
            doc_len_param = input_len
        else:
            doc_len_param = input_len - query_len - n_sep_tokens

774
775
776
777
778
779
780
        doc_lens, _, doc_offsets = get_sampling_params(
            self._rng,
            num_requests,
            range_ratio,
            doc_len_param,
            0,
            tokenizer,
781
        )
782

783
        vocab_size = tokenizer.vocab_size
784
785
786
        prohibited_tokens = tokenizer.all_special_ids
        all_tokens = np.arange(vocab_size)
        allowed_tokens = np.array(list(set(all_tokens) - set(prohibited_tokens)))
787
788
789
790
791
792
793
794
795
796

        query_prompt, query_input_len, token_mismatch_total = (
            self.generate_token_sequence(
                tokenizer=tokenizer,
                prefix_token_ids=[],
                prefix_len=0,
                vocab_size=vocab_size,
                input_len=query_len,
                offset=int(query_offsets[0]),
                index=0,
797
                allowed_tokens=allowed_tokens,
798
799
800
801
802
803
804
805
806
807
808
809
810
            )
        )

        requests = []
        for i in range(num_requests):
            prompt, total_input_len, token_mismatch = self.generate_token_sequence(  # noqa: E501
                tokenizer=tokenizer,
                prefix_token_ids=[],
                prefix_len=0,
                vocab_size=vocab_size,
                input_len=int(doc_lens[i]),
                offset=int(doc_offsets[i]),
                index=i + 1,
811
                allowed_tokens=allowed_tokens,
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
            )
            token_mismatch_total += token_mismatch
            requests.append((prompt, total_input_len))

        batch_requests = []
        # Create batched requests
        for i in range(0, num_requests, batchsize):
            batch = requests[i : i + batchsize]
            query_contrib = (
                (query_input_len + n_sep_tokens) * len(batch)
                if is_reranker
                else query_input_len
            )
            batch_requests.append(
                SampleRequest(
                    prompt=[query_prompt] + [req[0] for req in batch],
                    prompt_len=query_contrib + sum(req[1] for req in batch),
                    expected_output_len=0,
                    request_id=request_id_prefix + str(i // batchsize),
                )
            )

        if token_mismatch_total != 0:
            logger.warning(
                "Across all generated prompts, there were %d %s tokens "
                "than expected after decoding and re-encoding. This is "
                "expected due to the imperfect nature of the sampling "
                "procedure.",
                abs(token_mismatch_total),
                "more" if token_mismatch_total > 0 else "fewer",
            )

        return batch_requests


847
848
849
850
# -----------------------------------------------------------------------------
# MultiModalDataset Implementation
# -----------------------------------------------------------------------------

851

852
853
854
855
856
857
class RandomMultiModalDataset(RandomDataset):
    """
    Synthetic multimodal dataset (text + images) that extends RandomDataset.

    Status:
    - Images: supported via synthetic RGB data.
858
    - Video: supported via synthetic RGB data.
859
860
861
862
863
864
865
866
    - Audio: not yet supported.

    Sampling overview:
    1) Number of items per request is sampled uniformly from the integer range
       [floor(n·(1−r)), ceil(n·(1+r))], where n is the base count and r is
       `num_mm_items_range_ratio` in [0, 1]. r=0 keeps it fixed; r=1 allows 0.
       The maximum is further clamped to the sum of per-modality limits.
    2) Each item’s modality and shape is sampled from `bucket_config`, a dict
867
       mapping (height, width, num_frames) → probability. We treat
868
       `num_frames`=1 as image and `num_frames` > 1 as video.
869
       Entries with zero probability are removed and the rest are renormalized
870
871
872
873
874
875
876
       to sum to 1.
    3) Per-modality hard caps are enforced via `limit_mm_per_prompt`.
       When a modality reaches its cap, all of its buckets are excluded and the
       remaining probabilities are renormalized.

    Example bucket configuration:
    {(256, 256, 1): 0.5, (720, 1280, 1): 0.4, (720, 1280, 16): 0.1}
877
878
      - Two image buckets (`num_frames`=1) and one video bucket
      (`num_frames`=16).
879
880
881
882
    OBS.: Only image sampling is supported for now.
    """

    IS_MULTIMODAL = True
883
    DEFAULT_LIMIT_MM_PER_PROMPT = {"image": 255, "video": 1}
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898

    DEFAULT_BASE_ITEMS_PER_REQUEST = 1
    DEFAULT_NUM_MM_ITEMS_RANGE_RATIO = 0.0
    DEFAULT_MM_ITEM_BUCKET_CONFIG = {
        (256, 256, 1): 0.5,
        (720, 1280, 1): 0.5,
        (720, 1280, 16): 0.0,
    }
    DEFAULT_ENABLE_MULTIMODAL_CHAT = False

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

    def generate_synthetic_image(self, width: int, height: int) -> Image.Image:
        """Generate synthetic PIL image with random RGB values.
899
900
901

        NOTE: iid pixel sampling results in worst-case compression
        (good for stressing I/O), but very unlike real photos.
902
903
904
905
906
907
908
909
910
911
912
        We could consider a “low-freq” mode (e.g., noise blur)
        to emulate network realism instead of max stress.
        """
        random_pixels = self._rng.integers(
            0,
            256,
            (height, width, 3),
            dtype=np.uint8,
        )
        return Image.fromarray(random_pixels)

913
914
915
    def generate_synthetic_video(
        self, width: int, height: int, num_frames: int
    ) -> dict:
916
        """Generate synthetic video with random values.
917

918
919
        Creates a video with random pixel values, encodes it to MP4 format,
        and returns the content as bytes.
920
        """
921
922
        import cv2

923
924
925
926
927
928
929
930
931
932
933
        random_pixels = self._rng.integers(
            0,
            256,
            (num_frames, height, width, 3),
            dtype=np.uint8,
        )

        # Create a temporary video file in memory
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        fps = 30  # frames per second

934
        with NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
            temp_path = temp_file.name

            # Create video writer
            video_writer = cv2.VideoWriter(
                temp_path, fourcc=fourcc, fps=fps, frameSize=(width, height)
            )

            if not video_writer.isOpened():
                raise RuntimeError("Failed to create video writer")

            for frame in random_pixels:
                video_writer.write(frame)

            video_writer.release()
            temp_file.close()

            # Read the video file content
            with open(temp_path, "rb") as f:
                video_content = f.read()

            return {"bytes": video_content}
956
957
958
959
960
961
962
963
964
965

    def map_config_to_modality(self, config: tuple[int, int, int]) -> str:
        """Map the configuration to the modality."""
        if config[-1] == 1:
            return "image"
        elif config[-1] > 1:
            return "video"
        else:
            raise ValueError(f"Invalid multimodal item configuration: {config}")

966
967
968
    def normalize_bucket_config(
        self, bucket_config: dict[tuple[int, int, int], float]
    ) -> dict[tuple[int, int, int], float]:
969
970
971
972
973
974
975
976
977
978
979
        """
        Remove zero probability entries
        and normalize the bucket config to sum to 1.
        """
        # Raise error if value is negative
        if any(v < 0 for v in bucket_config.values()):
            raise ValueError("Bucket config values must be non-negative.")
        # Remove zero probability entries
        bucket_config = {k: v for k, v in bucket_config.items() if v > 0}
        # if bucket config is empty, raise error
        if not bucket_config:
980
981
982
            raise ValueError(
                "Got invalid bucket config. Bucket config values must be non-zero."
            )
983
984
985
986
        # Normalize the remaining bucket config to sum to 1
        total = sum(bucket_config.values())
        return {k: v / total for k, v in bucket_config.items()}

987
988
989
990
    def generate_mm_item(
        self,
        mm_item_config: tuple[int, int, int],
    ) -> Mapping[str, Any]:
991
        """
992
        Create synthetic images and videos and
993
994
995
996
        apply process_image/process_video respectively.
        This follows the OpenAI API chat completions
        https://github.com/openai/openai-python
        """
997

998
        if self.map_config_to_modality(mm_item_config) == "image":
999
1000
1001
            return process_image(
                self.generate_synthetic_image(mm_item_config[1], mm_item_config[0])
            )
1002
        elif self.map_config_to_modality(mm_item_config) == "video":
1003
1004
1005
1006
1007
            return process_video(
                self.generate_synthetic_video(
                    mm_item_config[1], mm_item_config[0], mm_item_config[2]
                )
            )
1008
        else:
1009
            raise ValueError(f"Invalid multimodal item configuration: {mm_item_config}")
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029

    def get_mm_item_sampling_params(
        self,
        base_items_per_request: int,
        num_mm_items_range_ratio: float,
        limit_mm_per_prompt: dict[str, int],
        bucket_config: dict[tuple[int, int, int], float],
    ) -> tuple[int, int, dict[str, int], dict[tuple[int, int, int], float]]:
        """
        Get the sampling parameters for the multimodal items.
        """
        # Enforce num_mm_items_range_ratio <= 1
        if not (0.0 <= num_mm_items_range_ratio <= 1.0):
            raise ValueError("num_mm_items_range_ratio must be in [0, 1].")

        # Ensure modalities to sample are in limit_mm_per_prompt
        for k, v in bucket_config.items():
            # get modality from bucket config
            modality = self.map_config_to_modality(k)
            if modality not in limit_mm_per_prompt:
1030
1031
1032
1033
1034
                raise ValueError(
                    f"Modality {modality} is not in "
                    f"limit_mm_per_prompt: "
                    f"{limit_mm_per_prompt.keys()}"
                )
1035

1036
        # Remove zero probability entries
1037
1038
1039
        # and normalize bucket config to sum to 1
        bucket_config = self.normalize_bucket_config(bucket_config)
        logger.info(
1040
1041
            "Normalized bucket config: %s",
            bucket_config,
1042
1043
        )
        # Only consider limit per prompt for modalities in bucket config
1044
        allowed_modalities = {self.map_config_to_modality(cfg) for cfg in bucket_config}
1045
        limit_mm_per_prompt = {
1046
1047
            k: v for k, v in limit_mm_per_prompt.items() if k in allowed_modalities
        }
1048
        if not limit_mm_per_prompt:
1049
            raise ValueError("No valid limits for modalities present in bucket_config.")
1050
1051

        logger.info(
1052
1053
            "Updated mm-limit-per-prompt: %s",
            limit_mm_per_prompt,
1054
1055
1056
1057
1058
        )

        # Get max and min num mm items and ensure
        # it is at most the sum of limit_mm_per_prompt for all modalities
        max_num_mm_items = min(
1059
            sum(limit_mm_per_prompt.values()),
1060
            math.ceil(base_items_per_request * (1 + num_mm_items_range_ratio)),
1061
1062
1063
        )
        # Ensure min num mm items is at least 0
        min_num_mm_items = max(
1064
            0, math.floor(base_items_per_request * (1 - num_mm_items_range_ratio))
1065
1066
1067
        )
        # Raise error if min num mm items is greater than max num mm items
        if min_num_mm_items > max_num_mm_items:
1068
1069
1070
1071
            raise ValueError(
                f"Min num mm items is greater than max mm items: "
                f"{min_num_mm_items} > {max_num_mm_items}"
            )
1072

1073
1074
        logger.info(
            "Sampling number of multimodal items from [%s, %s]",
1075
1076
            min_num_mm_items,
            max_num_mm_items,
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
        )

        return (
            min_num_mm_items,
            max_num_mm_items,
            limit_mm_per_prompt,
            bucket_config,
        )

    def get_mm_item_iterator(
        self,
        min_num_mm_items: int,
        max_num_mm_items: int,
        bucket_config: dict[tuple[int, int, int], float],
        limit_mm_per_prompt: dict[str, int],
1092
    ) -> Iterator[tuple[int, int, int]]:
1093
1094
1095
1096
1097
        """
        Iterator over the multimodal items for each request
        whose size is between min_num_mm_items and max_num_mm_items.

        Loop over the bucket config and sample a multimodal item.
1098
1099
        Loop until the number of multimodal items sampled is equal to
        request_num_mm_items or limit of multimodal items per prompt
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
        for all modalities is reached.

        Note:
        - This function operates on a per-request shallow copy of
          `bucket_config` (tuple->float). The original dict passed to
          `sample` is not mutated. If this ever changes, a test
          is implemented and will fail.
        """
        # Get the number of multimodal items to sample
        request_num_mm_items = int(
            self._rng.integers(min_num_mm_items, max_num_mm_items + 1)
1111
        )
1112
1113
1114
1115
        # If request_num_mm_items is 0, yield an empty iterator
        if request_num_mm_items == 0:
            return
        # Initialize modality counters
1116
        modality_counter = {self.map_config_to_modality(k): 0 for k in bucket_config}
1117
1118
1119
1120
1121
        # Copy the bucket config to avoid modifying the original
        bucket_config_copy = bucket_config.copy()
        # Loop over the number of multimodal items to sample
        while sum(modality_counter.values()) < request_num_mm_items:
            # Sample a multimodal item config
1122
1123
1124
            mm_item_config = self._rng.choice(
                list(bucket_config_copy.keys()), p=list(bucket_config_copy.values())
            )
1125
1126
1127
1128
            modality = self.map_config_to_modality(mm_item_config)
            # Check that modality count is less than limit per prompt
            if modality_counter[modality] < limit_mm_per_prompt[modality]:
                modality_counter[modality] += 1
1129
                yield (mm_item_config)
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
            else:
                # If the counter is greater than the limit per prompt
                # set all multimodal items of this modality to 0
                for k, v in bucket_config_copy.items():
                    if self.map_config_to_modality(k) == modality:
                        bucket_config_copy[k] = 0
                # If all configs are 0, break the loop
                # This should not happen as request_num_mm_items is at most
                # the sum of limit_mm_per_prompt for all modalities
                if all(v == 0 for v in bucket_config_copy.values()):
1140
1141
1142
                    logger.warning(
                        "Exhausted all multimodal items of modality %s", modality
                    )
1143
1144
                    break
                # Renormalize the bucket config
1145
                bucket_config_copy = self.normalize_bucket_config(bucket_config_copy)
1146
1147
1148

    def sample(
        self,
1149
        tokenizer: TokenizerLike,
1150
1151
        num_requests: int,
        request_id_prefix: str = "",
1152
        no_oversample: bool = False,
1153
        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
1154
        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
1155
1156
        input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
1157
        batchsize: int = 1,
1158
1159
1160
        limit_mm_per_prompt: dict[str, int] = DEFAULT_LIMIT_MM_PER_PROMPT,
        base_items_per_request: int = DEFAULT_BASE_ITEMS_PER_REQUEST,
        num_mm_items_range_ratio: float = DEFAULT_NUM_MM_ITEMS_RANGE_RATIO,
1161
1162
1163
        bucket_config: dict[
            tuple[int, int, int], float
        ] = DEFAULT_MM_ITEM_BUCKET_CONFIG,
1164
1165
1166
        enable_multimodal_chat: bool = DEFAULT_ENABLE_MULTIMODAL_CHAT,
        **kwargs,
    ) -> list[SampleRequest]:
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
        if batchsize != 1:
            raise NotImplementedError(
                "batchsize > 1 is not supported for RandomMultiModalDataset."
            )

        input_lens, output_lens, offsets = get_sampling_params(
            self._rng,
            num_requests,
            range_ratio,
            input_len,
            output_len,
            tokenizer,
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
        )

        (
            min_num_mm_items,
            max_num_mm_items,
            limit_mm_per_prompt,
            bucket_config,
        ) = self.get_mm_item_sampling_params(
            base_items_per_request,
            num_mm_items_range_ratio,
            limit_mm_per_prompt,
            bucket_config,
        )

        vocab_size = tokenizer.vocab_size
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
        # Can't use tokenizer.all_special_ids since
        # it returns ONLY ids from special_tokens_map.json
        # We want to exclude placeholder tokens and all
        # tokens that indicate start/end of image as it
        # may break prompt replacement logic.
        prohibited_tokens = list(
            tok_id
            for tok_id, token in tokenizer.added_tokens_decoder.items()
            if token.special
        )
        all_tokens = np.arange(vocab_size)
        allowed_tokens = np.array(list(set(all_tokens) - set(prohibited_tokens)))
        logger.debug(
            "Sampling from %d out of %d (vocab size)", len(allowed_tokens), vocab_size
        )
        # Generate prefix once
1210
        prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)
1211
1212
        # Add synthetic multimodal items to each request
        mm_requests = []
1213
        token_mismatch_total = 0
1214
        for i in range(num_requests):
1215
            prompt, total_input_len, token_mismatch = self.generate_token_sequence(  # noqa: E501
1216
1217
1218
1219
1220
1221
1222
                tokenizer=tokenizer,
                prefix_token_ids=prefix_token_ids,
                prefix_len=prefix_len,
                vocab_size=vocab_size,
                input_len=int(input_lens[i]),
                offset=int(offsets[i]),
                index=i,
1223
                allowed_tokens=allowed_tokens,
1224
            )
1225
            token_mismatch_total += token_mismatch
1226
1227
1228
1229
1230
1231
1232
1233
            # Get multimodal item iterator for a given request
            mm_item_iterator = self.get_mm_item_iterator(
                min_num_mm_items,
                max_num_mm_items,
                bucket_config,
                limit_mm_per_prompt,
            )

1234
1235
1236
1237
1238
1239
1240
            mm_content = cast(
                list[dict[str, Any]],
                [
                    self.generate_mm_item(mm_item_config)
                    for mm_item_config in mm_item_iterator
                ],
            )
1241
1242

            if enable_multimodal_chat:
1243
                # NOTE: For now this option is only provided for completeness
1244
1245
1246
                # given that the serve.py benchmark currently does not use it.
                mm_chat_prompt: Any = prompt
                mm_chat_prompt = self.apply_multimodal_chat_transformation(
1247
1248
                    prompt, mm_content
                )
1249
1250
1251
1252
1253
1254
1255
1256
1257
                sample_request = SampleRequest(
                    prompt=mm_chat_prompt,
                    prompt_len=total_input_len,
                    expected_output_len=int(output_lens[i]),
                    multi_modal_data=None,
                    request_id=request_id_prefix + str(i),
                )
            else:
                sample_request = SampleRequest(
1258
1259
1260
                    prompt=prompt,
                    prompt_len=total_input_len,
                    expected_output_len=int(output_lens[i]),
1261
                    multi_modal_data=mm_content,
1262
                    request_id=request_id_prefix + str(i),
1263
1264
                )
            mm_requests.append(sample_request)
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276

        if token_mismatch_total != 0:
            sign = "more" if token_mismatch_total > 0 else "fewer"
            logger.warning(
                "Across all generated prompts, there were %d %s tokens "
                "than expected after decoding and re-encoding. This is "
                "expected due to the imperfect nature of the sampling "
                "procedure.",
                abs(token_mismatch_total),
                sign,
            )

1277
        return mm_requests
1278

1279

1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
# -----------------------------------------------------------------------------
# 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 = [
1303
1304
            entry
            for entry in self.data
1305
1306
1307
            if "conversations" in entry and len(entry["conversations"]) >= 2
        ]
        random.seed(self.random_seed)
1308
1309
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(self.data)
1310
1311
1312

    def sample(
        self,
1313
        tokenizer: TokenizerLike,
1314
        num_requests: int,
1315
1316
        request_id_prefix: str = "",
        no_oversample: bool = False,
1317
1318
1319
        lora_path: str | None = None,
        max_loras: int | None = None,
        output_len: int | None = None,
1320
        enable_multimodal_chat: bool = False,
1321
        lora_assignment: str = "random",
1322
        **kwargs,
1323
1324
    ) -> list[SampleRequest]:
        samples: list[SampleRequest] = []
1325
        ind = 0
1326
1327
1328
1329
1330
1331
1332
1333
        for entry in self.data:
            if len(samples) >= num_requests:
                break
            prompt, completion = (
                entry["conversations"][0]["value"],
                entry["conversations"][1]["value"],
            )

1334
1335
1336
1337
1338
            lora_request = self.get_lora_request(
                index=ind,
                max_loras=max_loras,
                lora_path=lora_path,
                lora_assignment=lora_assignment,
1339
            )
1340
1341
1342
            prompt_ids = tokenizer(prompt).input_ids
            completion_ids = tokenizer(completion).input_ids
            prompt_len = len(prompt_ids)
1343
1344
1345
1346
1347
1348
            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,
            ):
1349
                continue
1350
1351
1352
            if image_path := entry.get("image"):
                mm_content = process_image(image_path)
            elif video_path := entry.get("video"):
1353
                mm_content = process_video(video_path)
1354
            else:
1355
                mm_content = None
1356
            if enable_multimodal_chat:
1357
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
1358
1359
1360
1361
1362
1363
            samples.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=new_output_len,
                    lora_request=lora_request,
1364
                    multi_modal_data=mm_content,
1365
                    request_id=request_id_prefix + str(ind),
1366
1367
                )
            )
1368
            ind += 1
1369
1370
1371
        self.maybe_oversample_requests(
            samples, num_requests, request_id_prefix, no_oversample
        )
1372
1373
1374
        return samples


1375
1376
class _ValidateDatasetArgs(argparse.Action):
    """Argparse action to validate dataset name and path compatibility."""
1377

1378
1379
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, values)
1380

1381
        # Get current values of both dataset_name and dataset_path
1382
1383
        dataset_name = getattr(namespace, "dataset_name", "random")
        dataset_path = getattr(namespace, "dataset_path", None)
1384

1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
        # Validate the combination
        if dataset_name == "random" and dataset_path is not None:
            parser.error(
                "Cannot use 'random' dataset with --dataset-path. "
                "Please specify the appropriate --dataset-name (e.g., "
                "'sharegpt', 'custom', 'sonnet') for your dataset file: "
                f"{dataset_path}"
            )


1395
def add_dataset_parser(parser: FlexibleArgumentParser):
1396
1397
1398
1399
1400
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="Trust remote code from huggingface",
    )
1401
1402
1403
1404
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--num-prompts",
        type=int,
1405
        default=DEFAULT_NUM_PROMPTS,
1406
1407
1408
1409
1410
1411
        help="Number of prompts to process.",
    )
    parser.add_argument(
        "--dataset-name",
        type=str,
        default="random",
1412
        action=_ValidateDatasetArgs,
1413
        choices=[
1414
1415
1416
1417
1418
            "sharegpt",
            "burstgpt",
            "sonnet",
            "random",
            "random-mm",
1419
            "random-rerank",
1420
1421
            "hf",
            "custom",
1422
            "custom_mm",
1423
1424
            "prefix_repetition",
            "spec_bench",
1425
        ],
1426
1427
        help="Name of the dataset to benchmark on.",
    )
1428
1429
1430
1431
1432
    parser.add_argument(
        "--no-stream",
        action="store_true",
        help="Do not load the dataset in streaming mode.",
    )
1433
1434
1435
1436
    parser.add_argument(
        "--dataset-path",
        type=str,
        default=None,
1437
        action=_ValidateDatasetArgs,
1438
1439
        help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
        "using HF dataset.",
1440
    )
1441
1442
1443
    parser.add_argument(
        "--no-oversample",
        action="store_true",
1444
        help="Do not oversample if the dataset has fewer samples than num-prompts.",
1445
    )
1446
1447
1448
    parser.add_argument(
        "--skip-chat-template",
        action="store_true",
1449
        help="Skip applying chat template to prompt for datasets that support it.",
1450
    )
1451
1452
1453
1454
1455
    parser.add_argument(
        "--enable-multimodal-chat",
        action="store_true",
        help="Enable multimodal chat transformation for datasets that support it.",
    )
1456
1457
1458
1459
1460
    parser.add_argument(
        "--disable-shuffle",
        action="store_true",
        help="Disable shuffling of dataset samples for deterministic ordering.",
    )
1461
1462
1463
1464
1465
1466
1467

    # group for dataset specific arguments
    custom_group = parser.add_argument_group("custom dataset options")
    custom_group.add_argument(
        "--custom-output-len",
        type=int,
        default=256,
1468
1469
1470
        help="Number of output tokens per request. Unless it is set to -1, the "
        "value overrides potential output length loaded from the dataset. It is "
        "used only for custom dataset.",
1471
1472
    )

1473
1474
1475
1476
1477
    spec_bench_group = parser.add_argument_group("spec bench dataset options")
    spec_bench_group.add_argument(
        "--spec-bench-output-len",
        type=int,
        default=256,
1478
        help="Num of output tokens per request, used only for spec bench dataset.",
1479
1480
1481
1482
1483
    )
    spec_bench_group.add_argument(
        "--spec-bench-category",
        type=str,
        default=None,
1484
        help="Category for spec bench dataset. If None, use all categories.",
1485
1486
    )

1487
1488
1489
1490
1491
    sonnet_group = parser.add_argument_group("sonnet dataset options")
    sonnet_group.add_argument(
        "--sonnet-input-len",
        type=int,
        default=550,
1492
        help="Number of input tokens per request, used only for sonnet dataset.",
1493
1494
1495
1496
1497
    )
    sonnet_group.add_argument(
        "--sonnet-output-len",
        type=int,
        default=150,
1498
        help="Number of output tokens per request, used only for sonnet dataset.",
1499
1500
1501
1502
1503
    )
    sonnet_group.add_argument(
        "--sonnet-prefix-len",
        type=int,
        default=200,
1504
        help="Number of prefix tokens per request, used only for sonnet dataset.",
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
    )

    sharegpt_group = parser.add_argument_group("sharegpt dataset options")
    sharegpt_group.add_argument(
        "--sharegpt-output-len",
        type=int,
        default=None,
        help="Output length for each request. Overrides the output length "
        "from the ShareGPT dataset.",
    )

1516
1517
1518
1519
1520
    blazedit_group = parser.add_argument_group("blazedit dataset options")
    blazedit_group.add_argument(
        "--blazedit-min-distance",
        type=float,
        default=0.0,
1521
        help="Minimum distance for blazedit dataset. Min: 0, Max: 1.0",
1522
1523
1524
1525
1526
    )
    blazedit_group.add_argument(
        "--blazedit-max-distance",
        type=float,
        default=1.0,
1527
        help="Maximum distance for blazedit dataset. Min: 0, Max: 1.0",
1528
1529
    )

1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
    asr_group = parser.add_argument_group("asr dataset options")
    asr_group.add_argument(
        "--asr-max-audio-len-sec",
        type=float,
        default=float("inf"),
        help="Maximum audio length in seconds for ASR dataset.",
    )
    asr_group.add_argument(
        "--asr-min-audio-len-sec",
        type=float,
        default=0.0,
        help="Minimum audio length in seconds for ASR dataset.",
    )

1544
    random_group = parser.add_argument_group("random dataset options")
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
    add_random_dataset_base_args(random_group)

    random_mm_group = parser.add_argument_group(
        "random multimodal dataset options extended from random dataset"
    )
    add_random_multimodal_dataset_args(random_mm_group)

    hf_group = parser.add_argument_group("hf dataset options")
    hf_group.add_argument(
        "--hf-subset", type=str, default=None, help="Subset of the HF dataset."
    )
    hf_group.add_argument(
        "--hf-split", type=str, default=None, help="Split of the HF dataset."
    )
    hf_group.add_argument(
        "--hf-name",
        type=str,
        default=None,
        help=(
            "Name of the dataset on HuggingFace "
            "(e.g., 'lmarena-ai/VisionArena-Chat'). "
            "Specify this if your dataset-path is a local path."
        ),
    )
    hf_group.add_argument(
        "--hf-output-len",
        type=int,
        default=None,
        help="Output length for each request. Overrides the output lengths "
        "from the sampled HF dataset.",
    )

    prefix_repetition_group = parser.add_argument_group(
        "prefix repetition dataset options"
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-prefix-len",
        type=int,
        default=256,
        help="Number of prefix tokens per request, used only for prefix "
        "repetition dataset.",
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-suffix-len",
        type=int,
        default=256,
        help="Number of suffix tokens per request, used only for prefix "
        "repetition dataset. Total input length is prefix_len + suffix_len.",
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-num-prefixes",
        type=int,
        default=10,
        help="Number of prefixes to generate, used only for prefix repetition "
        "dataset. Prompts per prefix is num_requests // num_prefixes.",
    )
    prefix_repetition_group.add_argument(
        "--prefix-repetition-output-len",
        type=int,
        default=128,
        help="Number of output tokens per request, used only for prefix "
        "repetition dataset.",
    )


def add_random_dataset_base_args(
    parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup,
) -> None:
    """Add CLI arguments for base random dataset options.

    This function adds arguments needed for:
    - random (random dataset)
    - random-mm (random multimodal dataset)
    - random-rerank (random dataset for reranking)

    Args:
        parser_or_group: Either a parser or an argument group to add arguments to.
    """
    parser_or_group.add_argument(
1624
1625
1626
        "--random-input-len",
        type=int,
        default=1024,
1627
        help="Number of input tokens per request, used only for random sampling.",
1628
    )
1629
    parser_or_group.add_argument(
1630
1631
1632
        "--random-output-len",
        type=int,
        default=128,
1633
        help="Number of output tokens per request, used only for random sampling.",
1634
    )
1635
    parser_or_group.add_argument(
1636
        "--random-range-ratio",
1637
1638
        type=str,
        default="0.0",
1639
        help="Range ratio for sampling input/output length, "
1640
1641
1642
        "used only for random sampling. A single float applies to both "
        'ISL and OSL. A JSON dict like \'{"input": 0.3, "output": 0.5}\' '
        "sets them independently. Values must be in [0, 1).",
1643
    )
1644
    parser_or_group.add_argument(
1645
1646
1647
        "--random-prefix-len",
        type=int,
        default=0,
1648
1649
1650
1651
1652
1653
1654
1655
        help=(
            "Number of fixed prefix tokens before the random context "
            "in a request. "
            "The total input length is the sum of `random-prefix-len` and "
            "a random "
            "context length sampled from [input_len * (1 - range_ratio), "
            "input_len * (1 + range_ratio)]."
        ),
1656
    )
1657
    parser_or_group.add_argument(
1658
1659
1660
        "--random-batch-size",
        type=int,
        default=1,
1661
        help=("Batch size for random sampling. Only used for embeddings benchmark."),
1662
    )
1663
    parser_or_group.add_argument(
1664
1665
1666
1667
1668
1669
1670
        "--no-reranker",
        action="store_true",
        help=(
            "Whether the model supports reranking natively."
            " Only used for reranker benchmark."
        ),
    )
1671

1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684

def add_random_multimodal_dataset_args(
    parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup,
) -> None:
    """Add CLI arguments for random multimodal dataset options.

    This function adds arguments needed for:
    - random-mm (random multimodal dataset)

    Args:
        parser_or_group: Either a parser or an argument group to add arguments to.
    """
    parser_or_group.add_argument(
1685
1686
1687
1688
1689
1690
1691
1692
1693
        "--random-mm-base-items-per-request",
        type=int,
        default=RandomMultiModalDataset.DEFAULT_BASE_ITEMS_PER_REQUEST,
        help=(
            "Base number of multimodal items per request for random-mm. "
            "Actual per-request count is sampled around this base using "
            "--random-mm-num-mm-items-range-ratio."
        ),
    )
1694
    parser_or_group.add_argument(
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
        "--random-mm-num-mm-items-range-ratio",
        type=float,
        default=RandomMultiModalDataset.DEFAULT_NUM_MM_ITEMS_RANGE_RATIO,
        help=(
            "Range ratio r in [0, 1] for sampling items per request. "
            "We sample uniformly from the closed integer range "
            "[floor(n*(1-r)), ceil(n*(1+r))] "
            "where n is the base items per request. "
            "r=0 keeps it fixed; r=1 allows 0 items. The maximum is clamped "
            "to the sum of per-modality limits from "
            "--random-mm-limit-mm-per-prompt. "
            "An error is raised if the computed min exceeds the max."
        ),
    )
1709
    parser_or_group.add_argument(
1710
1711
1712
1713
1714
        "--random-mm-limit-mm-per-prompt",
        type=json.loads,
        default=RandomMultiModalDataset.DEFAULT_LIMIT_MM_PER_PROMPT,
        help=(
            "Per-modality hard caps for items attached per request, e.g. "
1715
            '\'{"image": 3, "video": 0}\'. The sampled per-request item '
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
            "count is clamped to the sum of these limits. When a modality "
            "reaches its cap, its buckets are excluded and probabilities are "
            "renormalized."
            "OBS.: Only image sampling is supported for now."
        ),
    )

    def _parse_mm_bucket_config(v: object) -> dict[tuple[int, int, int], float]:
        # If already a dict (e.g., programmatic call), normalize keys
        def normalize(d: dict) -> dict[tuple[int, int, int], float]:
            out: dict[tuple[int, int, int], float] = {}
            for k, val in d.items():
                key = k
                if isinstance(key, str):
                    with suppress(Exception):
                        key = ast.literal_eval(key)
1732
1733
1734
1735
1736
                if not (
                    isinstance(key, tuple)
                    and len(key) == 3
                    and all(isinstance(x, int) for x in key)
                ):
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
                    raise ValueError(
                        f"Invalid bucket key {k!r}. Expected tuple (H, W, T)."
                    )
                out[(int(key[0]), int(key[1]), int(key[2]))] = float(val)
            return out

        if isinstance(v, dict):
            return normalize(v)
        if isinstance(v, str):
            # Python literal (supports tuple keys)
            parsed = ast.literal_eval(v)
            if not isinstance(parsed, dict):
                raise ValueError("Bucket config must parse to a dict.")
            return normalize(parsed)
        raise ValueError("Unsupported value for --random-mm-bucket-config.")

1753
    parser_or_group.add_argument(
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
        "--random-mm-bucket-config",
        type=_parse_mm_bucket_config,
        default=RandomMultiModalDataset.DEFAULT_MM_ITEM_BUCKET_CONFIG,
        help=(
            "The bucket config is a dictionary mapping a multimodal item"
            "sampling configuration to a probability."
            "Currently allows for 2 modalities: images and videos. "
            "An bucket key is a tuple of (height, width, num_frames)"
            "The value is the probability of sampling that specific item. "
            "Example: "
            "--random-mm-bucket-config "
            "{(256, 256, 1): 0.5, (720, 1280, 1): 0.4, (720, 1280, 16): 0.10} "
            "First item: images with resolution 256x256 w.p. 0.5"
            "Second item: images with resolution 720x1280 w.p. 0.4 "
            "Third item: videos with resolution 720x1280 and 16 frames w.p. 0.1"
            "OBS.: If the probabilities do not sum to 1, they are normalized."
            "OBS bis.: Only image sampling is supported for now."
        ),
1772
1773
    )

1774

1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
def _parse_range_ratio(value: str) -> RangeRatio:
    """Parse a ``--random-range-ratio`` CLI string.

    Accepts either a plain float (``"0.3"``) or a JSON dict
    (``'{"input": 0.3, "output": 0.5}'``).
    """
    try:
        return float(value)
    except ValueError:
        return json.loads(value)


1787
def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
1788
1789
1790
    if not hasattr(args, "request_id_prefix"):
        args.request_id_prefix = ""

1791
1792
1793
    if hasattr(args, "random_range_ratio") and isinstance(args.random_range_ratio, str):
        args.random_range_ratio = _parse_range_ratio(args.random_range_ratio)

1794
    if args.dataset_name == "custom":
1795
1796
1797
        dataset = CustomDataset(
            dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
        )
1798
1799
1800
1801
        input_requests = dataset.sample(
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
            output_len=args.custom_output_len,
1802
            skip_chat_template=args.skip_chat_template,
1803
            request_id_prefix=args.request_id_prefix,
1804
            no_oversample=args.no_oversample,
1805
1806
        )

1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
    elif args.dataset_name == "custom_mm":
        dataset = CustomMMDataset(
            dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
        )
        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,
        )

1820
    elif args.dataset_name == "sonnet":
1821
1822
1823
        dataset = SonnetDataset(
            dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
        )
1824
        # For the "sonnet" dataset, formatting depends on the backend.
1825
        if args.backend == "openai-chat":
1826
1827
1828
1829
1830
1831
1832
            input_requests = dataset.sample(
                num_requests=args.num_prompts,
                input_len=args.sonnet_input_len,
                output_len=args.sonnet_output_len,
                prefix_len=args.sonnet_prefix_len,
                tokenizer=tokenizer,
                return_prompt_formatted=False,
1833
                request_id_prefix=args.request_id_prefix,
1834
                no_oversample=args.no_oversample,
1835
1836
1837
            )
        else:
            assert tokenizer.chat_template or tokenizer.default_chat_template, (
1838
1839
                "Tokenizer/model must have chat template for sonnet dataset."
            )
1840
1841
1842
1843
1844
1845
1846
            input_requests = dataset.sample(
                num_requests=args.num_prompts,
                input_len=args.sonnet_input_len,
                output_len=args.sonnet_output_len,
                prefix_len=args.sonnet_prefix_len,
                tokenizer=tokenizer,
                return_prompt_formatted=True,
1847
                request_id_prefix=args.request_id_prefix,
1848
                no_oversample=args.no_oversample,
1849
1850
1851
1852
1853
            )

    elif args.dataset_name == "hf":
        # all following datasets are implemented from the
        # HuggingFaceDataset base class
1854
        hf_kwargs = {}
1855
1856
1857
1858
        if (
            args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in VisionArenaDataset.SUPPORTED_DATASET_PATHS
        ):
1859
            dataset_class = VisionArenaDataset
1860
            args.hf_split = args.hf_split if args.hf_split else "train"
1861
            args.hf_subset = None
1862
1863
1864
1865
1866
        elif (
            args.dataset_path in MMVUDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MMVUDataset.SUPPORTED_DATASET_PATHS
        ):
            dataset_class = MMVUDataset
1867
            args.hf_split = args.hf_split if args.hf_split else "validation"
1868
            args.hf_subset = None
1869
1870
1871
1872
        elif (
            args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in InstructCoderDataset.SUPPORTED_DATASET_PATHS
        ):
1873
            dataset_class = InstructCoderDataset
1874
            args.hf_split = args.hf_split if args.hf_split else "train"
1875
1876
1877
1878
        elif (
            args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MTBenchDataset.SUPPORTED_DATASET_PATHS
        ):
1879
            dataset_class = MTBenchDataset
1880
            args.hf_split = args.hf_split if args.hf_split else "train"
1881
1882
1883
1884
1885
        elif (
            args.dataset_path in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
        ):
            dataset_class = MultiModalConversationDataset
1886
1887
1888
1889
        elif (
            args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in ConversationDataset.SUPPORTED_DATASET_PATHS
        ):
1890
            dataset_class = ConversationDataset
1891
1892
1893
1894
        elif (
            args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in AIMODataset.SUPPORTED_DATASET_PATHS
        ):
1895
            dataset_class = AIMODataset
1896
            args.hf_split = args.hf_split if args.hf_split else "train"
1897
        elif (
1898
            args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS  # noqa: E501
1899
1900
            or args.hf_name in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS
        ):
1901
            dataset_class = NextEditPredictionDataset
1902
            args.hf_split = args.hf_split if args.hf_split else "train"
1903
1904
1905
1906
        elif (
            args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in ASRDataset.SUPPORTED_DATASET_PATHS
        ):
1907
            dataset_class = ASRDataset
1908
1909
1910
1911
1912
            args.hf_split = args.hf_split if args.hf_split else "train"
            hf_kwargs = {
                "asr_min_audio_len_sec": args.asr_min_audio_len_sec,
                "asr_max_audio_len_sec": args.asr_max_audio_len_sec,
            }
1913
1914
        elif args.dataset_path in BlazeditDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = BlazeditDataset
1915
            args.hf_split = args.hf_split if args.hf_split else "train"
1916
1917
1918
1919
            hf_kwargs = {
                "min_distance": args.blazedit_min_distance,
                "max_distance": args.blazedit_max_distance,
            }
1920
1921
1922
1923
        elif (
            args.dataset_path in MLPerfDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MLPerfDataset.SUPPORTED_DATASET_PATHS
        ):
1924
            dataset_class = MLPerfDataset
1925
            args.hf_split = args.hf_split if args.hf_split else "train"
1926
1927
1928
1929
1930
        elif (
            args.dataset_path in MMStarDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MMStarDataset.SUPPORTED_DATASET_PATHS
        ):
            dataset_class = MMStarDataset
1931
            args.hf_split = args.hf_split if args.hf_split else "val"
1932
            args.hf_subset = None
1933
        else:
1934
1935
1936
1937
1938
1939
1940
            supported_datasets = set(
                [
                    dataset_name
                    for cls in HuggingFaceDataset.__subclasses__()
                    for dataset_name in cls.SUPPORTED_DATASET_PATHS
                ]
            )
1941
1942
1943
1944
1945
            raise ValueError(
                f"Unsupported dataset path: {args.dataset_path}. "
                "Huggingface dataset only supports dataset_path"
                f" from one of following: {supported_datasets}. "
                "Please consider contributing if you would "
1946
1947
                "like to add support for additional dataset formats."
            )
1948

1949
1950
        if dataset_class.IS_MULTIMODAL and not (
            args.backend in ("openai-chat", "openai-audio")
1951
            or "embeddings-" in args.backend
1952
        ):
1953
1954
            # multi-modal benchmark is only available on OpenAI Chat
            # endpoint-type.
1955
1956
            raise ValueError(
                "Multi-modal content is only supported on 'openai-chat' and "
1957
1958
                "'openai-audio' backends."
            )
1959
1960
1961
1962
1963
        input_requests = dataset_class(
            dataset_path=args.dataset_path,
            dataset_subset=args.hf_subset,
            dataset_split=args.hf_split,
            random_seed=args.seed,
1964
            no_stream=args.no_stream,
1965
            hf_name=args.hf_name,
1966
            disable_shuffle=args.disable_shuffle,
1967
            trust_remote_code=args.trust_remote_code,
1968
1969
1970
1971
        ).sample(
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
            output_len=args.hf_output_len,
1972
            enable_multimodal_chat=args.enable_multimodal_chat,
1973
            request_id_prefix=args.request_id_prefix,
1974
            no_oversample=args.no_oversample,
1975
            skip_chat_template=args.skip_chat_template,
1976
            **hf_kwargs,
1977
1978
1979
1980
1981
        )

    else:
        # For datasets that follow a similar structure, use a mapping.
        dataset_mapping = {
1982
            "spec_bench": lambda: SpecBench(
1983
1984
1985
                dataset_path=args.dataset_path,
                category=args.spec_bench_category,
                disable_shuffle=args.disable_shuffle,
1986
            ).sample(
1987
1988
1989
                num_requests=args.num_prompts,
                tokenizer=tokenizer,
                output_len=args.spec_bench_output_len,
1990
                enable_multimodal_chat=args.enable_multimodal_chat,
1991
                request_id_prefix=args.request_id_prefix,
1992
                no_oversample=args.no_oversample,
1993
            ),
1994
            "sharegpt": lambda: ShareGPTDataset(
1995
1996
1997
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
1998
1999
2000
2001
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                output_len=args.sharegpt_output_len,
2002
                enable_multimodal_chat=args.enable_multimodal_chat,
2003
                request_id_prefix=args.request_id_prefix,
2004
                no_oversample=args.no_oversample,
2005
2006
            ),
            "burstgpt": lambda: BurstGPTDataset(
2007
2008
2009
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2010
2011
2012
2013
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                request_id_prefix=args.request_id_prefix,
2014
                no_oversample=args.no_oversample,
2015
2016
            ),
            "random": lambda: RandomDataset(
2017
2018
2019
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2020
            ).sample(
2021
2022
2023
2024
2025
2026
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                prefix_len=args.random_prefix_len,
                input_len=args.random_input_len,
                output_len=args.random_output_len,
                range_ratio=args.random_range_ratio,
2027
                request_id_prefix=args.request_id_prefix,
2028
                batchsize=args.random_batch_size,
2029
                no_oversample=args.no_oversample,
2030
            ),
2031
            "random-mm": lambda: RandomMultiModalDataset(
2032
2033
2034
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                prefix_len=args.random_prefix_len,
                range_ratio=args.random_range_ratio,
                input_len=args.random_input_len,
                output_len=args.random_output_len,
                base_items_per_request=args.random_mm_base_items_per_request,
                limit_mm_per_prompt=args.random_mm_limit_mm_per_prompt,
                num_mm_items_range_ratio=args.random_mm_num_mm_items_range_ratio,
                bucket_config=args.random_mm_bucket_config,
2046
                enable_multimodal_chat=args.enable_multimodal_chat,
2047
                request_id_prefix=args.request_id_prefix,
2048
                no_oversample=args.no_oversample,
2049
            ),
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
            "random-rerank": lambda: RandomDatasetForReranking(
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                input_len=args.random_input_len,
                range_ratio=args.random_range_ratio,
                request_id_prefix=args.request_id_prefix,
                batchsize=args.random_batch_size,
                is_reranker=not args.no_reranker,
            ),
2063
            "prefix_repetition": lambda: PrefixRepetitionRandomDataset(
2064
2065
2066
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2067
2068
2069
2070
2071
2072
2073
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                prefix_len=args.prefix_repetition_prefix_len,
                suffix_len=args.prefix_repetition_suffix_len,
                num_prefixes=args.prefix_repetition_num_prefixes,
                output_len=args.prefix_repetition_output_len,
2074
                request_id_prefix=args.request_id_prefix,
2075
                no_oversample=args.no_oversample,
2076
            ),
2077
2078
2079
        }

        try:
2080
            # Enforce endpoint compatibility for multimodal datasets.
2081
            if args.dataset_name == "random-mm" and args.backend not in ["openai-chat"]:
2082
2083
2084
2085
                raise ValueError(
                    "Multi-modal content (images) is only supported on "
                    "'openai-chat' backend."
                )
2086
2087
2088
2089
2090
2091
2092
            input_requests = dataset_mapping[args.dataset_name]()
        except KeyError as err:
            raise ValueError(f"Unknown dataset: {args.dataset_name}") from err

    return input_requests


2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
# -----------------------------------------------------------------------------
# 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.,
    ```
2103
2104
2105
    {"prompt": "What is the capital of India?", "output_tokens": 10}
    {"prompt": "What is the capital of Iran?", "output_tokens": 1520}
    {"prompt": "What is the capital of China?", "output_tokens": 819}
2106
    ```
2107
2108
    Note that 'output_tokens' column is optional and has to be provided only if
    'custom-output-len' argument is None or -1.
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
    """

    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
2124
        self.data: list[dict] = []
2125
2126
2127

        # Load the JSONL file
        if self.dataset_path.endswith(".jsonl"):
2128
            jsonl_data = pd.read_json(path_or_buf=self.dataset_path, lines=True)
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141

            # 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(
2142
2143
                "Only JSONL format is supported for CustomDataset."
            )
2144
2145

        random.seed(self.random_seed)
2146
2147
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(self.data)
2148
2149
2150

    def sample(
        self,
2151
        tokenizer: TokenizerLike,
2152
        num_requests: int,
2153
2154
        request_id_prefix: str = "",
        no_oversample: bool = False,
2155
2156
2157
        lora_path: str | None = None,
        max_loras: int | None = None,
        output_len: int | None = None,
2158
2159
2160
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
        **kwargs,
2161
    ) -> list[SampleRequest]:
2162
2163
2164
2165
        # load all data if needed
        self.num_available_samples = len(self.data)
        if num_requests <= 0:
            num_requests = self.num_available_samples
2166
2167
2168
2169
2170
            logger.info(
                "num_requests is set to 0 or negative, "
                "so using all available samples: %d",
                num_requests,
            )
2171

2172
        sampled_requests: list[SampleRequest] = []
2173
        for i, item in enumerate(self.data):
2174
2175
2176
2177
            if len(sampled_requests) >= num_requests:
                break
            prompt = item["prompt"]

2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
            if tokenizer is None:
                new_output_len = 1
            else:
                new_output_len = output_len
                if output_len is None or output_len == -1:
                    # check that the request has an 'output_tokens' field
                    if "output_tokens" not in item:
                        raise ValueError(
                            "If no output length is provided the "
                            "custom dataset must contain an 'output_tokens' field."
                        )
                    # Use number of output tokens from the request data
                    try:
                        new_output_len = int(item["output_tokens"])
                    except (ValueError, TypeError) as e:
                        raise ValueError(
                            f"Invalid value for 'output_tokens' in custom dataset: "
                            f"'{item['output_tokens']}'. Must be an integer."
                        ) from e

            if tokenizer is None:
                prompt_len = 1
            else:
                # apply template
                if not skip_chat_template:
                    prompt = tokenizer.apply_chat_template(
                        [{"role": "user", "content": prompt}],
                        add_generation_prompt=True,
                        tokenize=False,
2207
                    )
2208

2209
                prompt_len = len(tokenizer(prompt).input_ids)
2210
2211
2212
2213
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
2214
                    expected_output_len=new_output_len,
2215
                    request_id=request_id_prefix + str(i),
2216
2217
2218
2219
2220
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2221
2222
2223
2224

        return sampled_requests


2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
class CustomMMDataset(CustomDataset):
    """
    Implements the Custom MultiModal dataset. Loads data from a JSONL file and generates
    sample requests based on conversation turns. E.g.,
    ```
    {
        "prompt": "How many red blocks in the given images?",
        "image_files": ["path/to/image1.png", "path/to/image2.png"],
    }
    {
        "prompt": "Which country has the most pokemons based on the given graphs?",
        "image_files": ["path/to/image.png"],
    }
    ```

    NOTE: Only the first image file in "image_files" is used for each sample request.

    This is used to benchmark multimodal LLMs on arbitrary datasets.
    """

    IS_MULTIMODAL = True

    def sample(
        self,
        tokenizer: TokenizerLike,
        num_requests: int,
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
        request_id_prefix: str = "",
        no_oversample: bool = False,
        **kwargs,
2256
    ) -> list[SampleRequest]:
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
        # load all data if needed
        self.num_available_samples = len(self.data)
        if num_requests <= 0:
            num_requests = self.num_available_samples
            logger.info(
                "num_requests is set to 0 or negative, "
                "so using all available samples: %d",
                num_requests,
            )

        sampled_requests = []
        for i, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
                break
            prompt = item["prompt"]

            prompt_len = len(tokenizer(prompt).input_ids)
            images = item["image_files"]
            if len(images) > 1:
                logger.warning(
                    "Multiple image files found for sample %d. "
                    "Only the first image will be used.",
                    i,
                )
            mm_content = process_image(images[0])
            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
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)

            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=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


2304
2305
2306
2307
2308
2309
2310
2311
# -----------------------------------------------------------------------------
# Spec Bench Dataset Implementation
# -----------------------------------------------------------------------------


class SpecBench(CustomDataset):
    """
    Implements the SpecBench dataset: https://github.com/hemingkx/Spec-Bench
2312
    Download the dataset using:
2313
    wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl
2314
    """  # noqa: E501
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327

    def __init__(self, **kwargs) -> None:
        self.category = kwargs.pop("category", 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 = []

        # Load the JSONL file
2328
        jsonl_data = pd.read_json(path_or_buf=self.dataset_path, lines=True)
2329
2330
2331
2332
2333
2334
2335

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

        for _, row in jsonl_data.iterrows():
            # sample only from a specific category if specified
2336
            if (not self.category) or (self.category == row["category"]):
2337
2338
2339
2340
                prompt = row["turns"][0]
                self.data.append({"prompt": prompt})

        random.seed(self.random_seed)
2341
2342
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(self.data)
2343

2344
2345
2346
    def sample(
        **kwargs,
    ) -> list[SampleRequest]:
2347
        # leverage CustomDataset sample
2348
2349
2350
        return super().sample(
            **kwargs,
        )
2351
2352


2353
2354
2355
2356
# -----------------------------------------------------------------------------
# Sonnet Dataset Implementation
# -----------------------------------------------------------------------------

2357

2358
2359
2360
@deprecated(
    "SonnetDataset is deprecated and will be removed in a future version.",
)
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
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()

    def sample(
        self,
2387
        tokenizer: TokenizerLike,
2388
        num_requests: int,
2389
2390
        request_id_prefix: str = "",
        no_oversample: bool = False,
2391
2392
2393
2394
2395
        prefix_len: int = DEFAULT_PREFIX_LEN,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        return_prompt_formatted: bool = False,
        **kwargs,
2396
    ) -> list[SampleRequest]:
2397
2398
        # Calculate average token length for a poem line.
        tokenized_lines = [tokenizer(line).input_ids for line in self.data]
2399
        avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)
2400
2401
2402
2403

        # 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}]
2404
2405
2406
        base_fmt = tokenizer.apply_chat_template(
            base_msg, add_generation_prompt=True, tokenize=False
        )
2407
2408
2409
2410
        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 "
2411
2412
                f"({base_offset})."
            )
2413
2414
2415
2416
2417
2418

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

2419
        samples: list[SampleRequest] = []
2420
        ind = 0
2421
        while len(samples) < num_requests:
2422
2423
2424
            extra_lines = random.choices(
                self.data, k=num_input_lines - num_prefix_lines
            )
2425
2426
2427
            prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"
            msg = [{"role": "user", "content": prompt}]
            prompt_formatted = tokenizer.apply_chat_template(
2428
2429
                msg, add_generation_prompt=True, tokenize=False
            )
2430
2431
2432
2433
            prompt_len = len(tokenizer(prompt_formatted).input_ids)
            if prompt_len <= input_len:
                samples.append(
                    SampleRequest(
2434
                        prompt=prompt_formatted if return_prompt_formatted else prompt,
2435
2436
                        prompt_len=prompt_len,
                        expected_output_len=output_len,
2437
2438
2439
                        request_id=request_id_prefix + str(ind),
                    )
                )
2440
                ind += 1
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
        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()

2460
2461
2462
    def load_data(
        self,
    ):
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
        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):
2476
            data = self.data.sample(n=num_requests, random_state=self.random_seed)
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
        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()

    def sample(
        self,
2488
        tokenizer: TokenizerLike,
2489
        num_requests: int,
2490
        request_id_prefix: str = "",
2491
        no_oversample: bool = False,
2492
        lora_assignment: str = "random",
2493
2494
        max_loras: int | None = None,
        lora_path: str | None = None,
2495
2496
2497
2498
2499
2500
2501
        **kwargs,
    ) -> list[SampleRequest]:
        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])
2502
2503
2504
2505
2506
            lora_req = self.get_lora_request(
                index=i,
                max_loras=max_loras,
                lora_path=lora_path,
                lora_assignment=lora_assignment,
2507
            )
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
            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,
2519
                    request_id=request_id_prefix + str(i),
2520
2521
                )
            )
2522
2523
2524
2525
2526
2527
2528
2529
2530
        return samples


# -----------------------------------------------------------------------------
# HuggingFace Dataset Base Implementation
# -----------------------------------------------------------------------------
class HuggingFaceDataset(BenchmarkDataset):
    """Base class for datasets hosted on HuggingFace."""

2531
    SUPPORTED_DATASET_PATHS: set[str] | dict[str, Callable] = set()
2532
2533
2534
2535
2536

    def __init__(
        self,
        dataset_path: str,
        dataset_split: str,
2537
        no_stream: bool = False,
2538
2539
        dataset_subset: str | None = None,
        hf_name: str | None = None,
2540
        trust_remote_code: bool = False,
2541
2542
2543
2544
2545
2546
        **kwargs,
    ) -> None:
        super().__init__(dataset_path=dataset_path, **kwargs)

        self.dataset_split = dataset_split
        self.dataset_subset = dataset_subset
2547
        self.load_stream = not no_stream
2548
        self.hf_name = hf_name or dataset_path
2549
        self.trust_remote_code = trust_remote_code
2550
2551
2552
2553
2554
2555
2556
2557
        self.load_data()

    def load_data(self) -> None:
        """Load data from HuggingFace datasets."""
        self.data = load_dataset(
            self.dataset_path,
            name=self.dataset_subset,
            split=self.dataset_split,
2558
            streaming=self.load_stream,
2559
            trust_remote_code=self.trust_remote_code,
2560
        )
2561
2562
        if not getattr(self, "disable_shuffle", False):
            self.data = self.data.shuffle(seed=self.random_seed)
2563
2564
2565
2566
2567
2568
2569
2570


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


class ConversationDataset(HuggingFaceDataset):
2571
    """Dataset for text-only conversation data."""
2572

2573
    SUPPORTED_DATASET_PATHS = {
2574
        "Aeala/ShareGPT_Vicuna_unfiltered",
2575
    }
2576
2577
2578
2579
    IS_MULTIMODAL = False

    def sample(
        self,
2580
        tokenizer: TokenizerLike,
2581
2582
2583
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
2584
2585
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2586
        **kwargs,
2587
    ) -> list[SampleRequest]:
2588
2589
        # Filter examples with at least 2 conversations
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
2590
        sampled_requests: list[SampleRequest] = []
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
        ind = 0
        dynamic_output = output_len is None

        for item in filtered_data:
            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
            if dynamic_output and not is_valid_sequence(prompt_len, completion_len):
                continue
            mm_content = process_image(item["image"]) if "image" in item else None
            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
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
                    request_id=request_id_prefix + str(ind),
                )
            )
            ind += 1
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
        return sampled_requests


class MultiModalConversationDataset(HuggingFaceDataset):
    """Dataset for multimodal conversation data."""

    SUPPORTED_DATASET_PATHS = {
        "lmms-lab/LLaVA-OneVision-Data",
    }
2636
    IS_MULTIMODAL = True
2637

2638
2639
    def sample(
        self,
2640
        tokenizer: TokenizerLike,
2641
2642
2643
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
2644
2645
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2646
        **kwargs,
2647
    ) -> list[SampleRequest]:
2648
        # Filter examples with at least 2 conversations
2649
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
2650
        sampled_requests: list[SampleRequest] = []
2651
        ind = 0
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
        dynamic_output = output_len is None

        for item in filtered_data:
            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
2666
            if dynamic_output and not is_valid_sequence(prompt_len, completion_len):
2667
                continue
2668
            mm_content = process_image(item["image"]) if "image" in item else None
2669
2670
2671
2672
            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
2673
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
2674
2675
2676
2677
2678
2679
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
2680
                    request_id=request_id_prefix + str(ind),
2681
2682
                )
            )
2683
            ind += 1
2684
2685
2686
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
        return sampled_requests


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


class VisionArenaDataset(HuggingFaceDataset):
    """
    Vision Arena Dataset.
    """

    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {
2702
2703
        "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"],
2704
    }
2705
    IS_MULTIMODAL = True
2706
2707
2708

    def sample(
        self,
2709
        tokenizer: TokenizerLike,
2710
        num_requests: int,
2711
        request_id_prefix: str = "",
2712
        no_oversample: bool = False,
2713
2714
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2715
        **kwargs,
2716
    ) -> list[SampleRequest]:
2717
2718
2719
2720
        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
        if parser_fn is None:
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")

2721
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2722

2723
        sampled_requests = []
2724
        for i, item in enumerate(self.data):
2725
2726
            if len(sampled_requests) >= num_requests:
                break
2727

2728
2729
            prompt = parser_fn(item)
            mm_content = process_image(item["images"][0])
2730
            prompt_len = len(tokenizer.encode(prompt))
2731
2732
2733
2734
            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
2735
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
2736

2737
2738
2739
2740
2741
2742
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
2743
                    request_id=request_id_prefix + str(i),
2744
2745
                )
            )
2746

2747
2748
2749
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2750
2751
2752
        return sampled_requests


2753
2754
2755
2756
2757
2758
2759
2760
class MMVUDataset(HuggingFaceDataset):
    """
    MMVU Dataset.
    https://huggingface.co/datasets/yale-nlp/MMVU
    """

    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {
2761
2762
2763
2764
2765
        "yale-nlp/MMVU": lambda x: (
            x["question"]
            + " "
            + (" ".join(f"{k}.{v}" for k, v in x["choices"].items()))
        ),
2766
2767
    }

2768
2769
2770
2771
2772
2773
2774
2775
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        self._remote_path_root = (
            f"https://huggingface.co/datasets/{self.hf_name}/resolve/main"
        )
        self._local_path_root = snapshot_download(self.hf_name, repo_type="dataset")

2776
2777
    def sample(
        self,
2778
        tokenizer: TokenizerLike,
2779
2780
2781
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
2782
2783
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2784
        **kwargs,
2785
    ) -> list[SampleRequest]:
2786
2787
2788
2789
        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
        if parser_fn is None:
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")

2790
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2791

2792
2793
2794
2795
        sampled_requests = []
        for i, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
                break
2796

2797
            prompt = parser_fn(item)
2798
2799
2800
            mm_content = process_video(
                item["video"].replace(self._remote_path_root, self._local_path_root)
            )
2801
            prompt_len = len(tokenizer.encode(prompt))
2802
2803
2804
2805
            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
2806
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
2807

2808
2809
2810
2811
2812
2813
2814
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
                    request_id=request_id_prefix + str(i),
2815
2816
                )
            )
2817

2818
2819
2820
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2821
2822
2823
        return sampled_requests


2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
# -----------------------------------------------------------------------------
# Instruct Coder Dataset Implementation
# -----------------------------------------------------------------------------


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

    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.
    """

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

2844
2845
    def sample(
        self,
2846
        tokenizer: TokenizerLike,
2847
        num_requests: int,
2848
2849
        request_id_prefix: str = "",
        no_oversample: bool = False,
2850
        output_len: int | None = None,
2851
2852
2853
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
        **kwargs,
2854
    ) -> list[SampleRequest]:
2855
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2856
        sampled_requests: list[SampleRequest] = []
2857
        for i, prompt in enumerate(self.sample_prompts(n=num_requests)):
2858
            # apply template
2859
2860
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
2861
                    [{"role": "user", "content": prompt}],
2862
2863
2864
                    add_generation_prompt=True,
                    tokenize=False,
                )
2865

2866
2867
2868
2869
2870
2871
            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
2872
                    request_id=request_id_prefix + str(i),
2873
2874
2875
2876
2877
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2878
2879
        return sampled_requests

2880
2881
2882
2883
2884
2885
2886
2887
    def sample_prompts(self, n: int) -> Iterator[str]:
        for item in self.data.take(n):
            prompt = (
                f"{item['input']}\n\n{item['instruction']} Just output "
                "the code, do not include any explanation."
            )
            yield prompt

2888

2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
# -----------------------------------------------------------------------------
# MT-Bench Dataset Implementation
# -----------------------------------------------------------------------------


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

    We create a single turn dataset for MT-Bench.
    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
    """  # noqa: E501

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

    def sample(
        self,
2911
        tokenizer: TokenizerLike,
2912
        num_requests: int,
2913
2914
        request_id_prefix: str = "",
        no_oversample: bool = False,
2915
        output_len: int | None = None,
2916
        enable_multimodal_chat: bool = False,
2917
        skip_chat_template: bool = False,
2918
        **kwargs,
2919
    ) -> list[SampleRequest]:
2920
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2921
        sampled_requests: list[SampleRequest] = []
2922

2923
        for i, item in enumerate(self.data):
2924
2925
2926
2927
2928
            if len(sampled_requests) >= num_requests:
                break
            prompt = item["turns"][0]

            # apply template
2929
2930
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
2931
                    [{"role": "user", "content": prompt}],
2932
2933
2934
                    add_generation_prompt=True,
                    tokenize=False,
                )
2935
2936
2937
2938
2939
2940
2941

            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
2942
                    request_id=request_id_prefix + str(i),
2943
2944
2945
2946
2947
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2948
2949
2950
        return sampled_requests


2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
# -----------------------------------------------------------------------------
# Blazedit Dataset Implementation
# -----------------------------------------------------------------------------


class BlazeditDataset(HuggingFaceDataset):
    """
    Blazedit Dataset.
    https://github.com/ise-uiuc/blazedit

    5k char version: vdaita/edit_5k_char
    10k char version: vdaita/edit_10k_char
    """  # noqa: E501

    # 5k char version will have output as ~5k chars
    # 10k char version will have output as ~10k chars
    # Assuming 3 char per token, 10k chars will be 3333 tokens
    # We set default to 4000 to be safe
    DEFAULT_OUTPUT_LEN = 4000
    SUPPORTED_DATASET_PATHS = {
        "vdaita/edit_5k_char",
        "vdaita/edit_10k_char",
    }

    def sample(
        self,
2977
        tokenizer: TokenizerLike,
2978
        num_requests: int,
2979
        output_len: int | None = None,
2980
        skip_chat_template: bool = False,
2981
        request_id_prefix: str = "",
2982
        no_oversample: bool = False,
2983
2984
2985
        min_distance: float = 0.0,
        max_distance: float = 1.0,
        **kwargs,
2986
    ) -> list[SampleRequest]:
2987
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
        sampled_requests = []

        for i, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
                break
            code = item["code"]
            change_request = item["change_request"]
            norm_distance = item["norm_distance"]

            # compare the levenshtein distance normalized by code length
            if norm_distance < min_distance or norm_distance > max_distance:
                continue
3000
3001

            # template copied from
3002
            # https://github.com/ise-uiuc/blazedit/blob/7765137e656fd62de877422d2e4cf8de51228054/dataset/create_refined_dataset.py#L94-L105 # noqa: E501
3003
            prompt = f"""Given a code file, please apply the change requests and generate the new file.
3004
3005
3006
3007
3008
3009
3010
3011
3012

Original file:
```python
{code}
```

Change request:
{change_request}

3013
Please generate the new code file in the "New file" section below."""  # noqa: E501
3014
3015

            # apply template
3016
3017
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
3018
                    [{"role": "user", "content": prompt}],
3019
3020
3021
                    add_generation_prompt=True,
                    tokenize=False,
                )
3022
3023
3024
3025
3026
3027
3028
3029
3030

            prompt_len = len(tokenizer(prompt).input_ids)

            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    request_id=request_id_prefix + str(i),
3031
3032
3033
3034
3035
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3036

3037
3038
3039
        return sampled_requests


3040
3041
3042
3043
3044
3045
3046
3047
3048
# -----------------------------------------------------------------------------
# AIMO Dataset Implementation
# -----------------------------------------------------------------------------


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

3050
    SUPPORTED_DATASET_PATHS = {
3051
3052
3053
        "AI-MO/aimo-validation-aime",
        "AI-MO/NuminaMath-1.5",
        "AI-MO/NuminaMath-CoT",
3054
3055
    }

3056
3057
    def sample(
        self,
3058
        tokenizer: TokenizerLike,
3059
3060
3061
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
3062
        output_len: int | None = None,
3063
        **kwargs,
3064
3065
    ) -> list[SampleRequest]:
        sampled_requests: list[SampleRequest] = []
3066
        ind = 0
3067
3068
3069
3070
3071
        dynamic_output = output_len is None

        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break
3072
            prompt, completion = item["problem"], item["solution"]
3073
3074
3075
3076
3077
3078
3079

            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
3080
3081
3082
            if dynamic_output and not is_valid_sequence(
                prompt_len, completion_len, max_prompt_len=2048, max_total_len=32000
            ):
3083
3084
3085
3086
3087
3088
3089
                continue
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=None,
3090
                    request_id=request_id_prefix + str(ind),
3091
3092
                )
            )
3093
            ind += 1
3094
3095
3096
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3097
        return sampled_requests
3098
3099
3100
3101
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117


# -----------------------------------------------------------------------------
# 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:

3118
"""  # noqa: E501
3119
3120
3121


def _format_zeta_prompt(
3122
3123
    sample: dict, original_start_marker: str = "<|editable_region_start|>"
) -> dict:
3124
    """Format the zeta prompt for the Next Edit Prediction (NEP) dataset.
3125
3126
3127

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

3130
    Args:
3131
        sample: The dataset sample containing events,
3132
            inputs, and outputs.
3133
3134
        original_start_marker: The marker indicating the
            start of the editable region. Defaults to
3135
            "<|editable_region_start|>".
3136

3137
3138
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
    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,
    }

3166
3167
    def sample(
        self,
3168
        tokenizer: TokenizerLike,
3169
3170
3171
3172
3173
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
        **kwargs,
    ):
3174
        formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(self.hf_name)
3175
        if formatting_prompt_func is None:
3176
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")
3177
        samples = []
3178
        for i, sample in enumerate(self.data):
3179
3180
3181
3182
3183
3184
            sample = formatting_prompt_func(sample)
            samples.append(
                SampleRequest(
                    prompt=sample["prompt"],
                    prompt_len=len(tokenizer(sample["prompt"]).input_ids),
                    expected_output_len=len(
3185
3186
                        tokenizer(sample["expected_output"]).input_ids
                    ),
3187
                    request_id=request_id_prefix + str(i),
3188
3189
                )
            )
3190
3191
            if len(samples) >= num_requests:
                break
3192
3193
3194
        self.maybe_oversample_requests(
            samples, num_requests, request_id_prefix, no_oversample
        )
3195
        return samples
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230


# -----------------------------------------------------------------------------
# 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                    |
    +----------------+----------------------------------------+--------------------------+-----------------------------+

    """  # noqa: E501

    SUPPORTED_DATASET_PATHS = {
        "openslr/librispeech_asr",
        "facebook/voxpopuli",
        "LIUM/tedlium",
        "edinburghcstr/ami",
        "speechcolab/gigaspeech",
        "kensho/spgispeech",
    }

3231
    DEFAULT_OUTPUT_LEN = 1024
3232
3233
3234
3235
    IS_MULTIMODAL = True

    def sample(
        self,
3236
        tokenizer: TokenizerLike,
3237
        num_requests: int,
3238
        request_id_prefix: str = "",
3239
        no_oversample: bool = False,
3240
        output_len: int | None = None,
3241
        **kwargs,
3242
    ) -> list[SampleRequest]:
3243
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
Ekagra Ranjan's avatar
Ekagra Ranjan committed
3244
        if "openai" in getattr(tokenizer, "name_or_path", ""):
3245
3246
3247
            prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
        else:
            prompt = ""
3248
        prompt_len = len(tokenizer(prompt).input_ids)
3249
        sampled_requests: list[SampleRequest] = []
3250
        ind = 0
3251
        skipped = 0
3252
3253
3254
        asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec")
        asr_max_audio_len_sec = kwargs.get("asr_max_audio_len_sec")
        durations = []
3255
3256
3257
3258
3259
        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break
            audio = item["audio"]
            y, sr = audio["array"], audio["sampling_rate"]
3260
            duration_s = get_audio_duration(y=y, sr=sr)
3261
            if duration_s < asr_min_audio_len_sec or duration_s > asr_max_audio_len_sec:
3262
3263
3264
                skipped += 1
                continue

3265
            durations.append(duration_s)
3266
3267
3268
3269
3270
3271
3272
            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,
3273
                    request_id=request_id_prefix + str(ind),
3274
3275
                )
            )
3276
            ind += 1
3277
3278
3279
3280
3281
3282
3283
        if skipped:
            logger.warning(
                "%d samples discarded from dataset due to"
                " their length being greater than"
                " what Whisper supports.",
                skipped,
            )
3284
3285
3286
3287
3288
3289
3290
3291
3292
3293
3294
3295
3296
3297

        logger.info("Number of audio samples: %d", len(durations))
        avg_duration = sum(durations) / len(durations) if durations else 0
        min_duration = min(durations) if durations else 0
        max_duration = max(durations) if durations else 0
        median_duration = np.median(durations) if durations else 0
        logger.info(
            "Audio duration statistics (s): avg=%.2f, min=%.2f, max=%.2f, median=%.2f",
            avg_duration,
            min_duration,
            max_duration,
            median_duration,
        )

3298
3299
3300
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3301
        return sampled_requests
3302
3303
3304
3305
3306
3307
3308
3309
3310
3311
3312
3313
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
3328
3329
3330
3331
3332
3333


# -----------------------------------------------------------------------------
# MLPerf Dataset Implementation
# -----------------------------------------------------------------------------


class MLPerfDataset(HuggingFaceDataset):
    """
    MLPerf Inference Dataset.

    Dataset on HF:
    https://huggingface.co/datasets/mgoin/mlperf-inference-llama2-data
    https://huggingface.co/datasets/mgoin/mlperf-inference-llama3.1-data

    Each record contains:
      - "system_prompt": system role instruction.
      - "question": user question.
      - "output": reference answer.

    We combine the system prompt and question into a chat-formatted prompt
    (using the tokenizer's chat template) and set the expected output length to
    the tokenized length of the provided reference answer.
    """

    SUPPORTED_DATASET_PATHS = {
        "mgoin/mlperf-inference-llama2-data",
        "mgoin/mlperf-inference-llama3.1-data",
    }

    def sample(
        self,
3334
        tokenizer: TokenizerLike,
3335
        num_requests: int,
3336
        request_id_prefix: str = "",
3337
        no_oversample: bool = False,
3338
        output_len: int | None = None,
3339
3340
3341
3342
3343
        **kwargs,
    ) -> list[SampleRequest]:
        # Force dynamic output length based on reference completion.
        dynamic_output = output_len is None
        sampled_requests: list[SampleRequest] = []
3344
        ind = 0
3345
3346
3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378

        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break

            system_prompt = item["system_prompt"]
            question = item["question"]
            reference_answer = item["output"]

            # Build chat-style prompt using tokenizer template, if available.
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": question},
            ]
            prompt_formatted = tokenizer.apply_chat_template(
                messages, add_generation_prompt=True, tokenize=False
            )
            prompt_len = len(tokenizer(prompt_formatted).input_ids)

            # Determine output length from reference answer tokens.
            ref_out_len = len(
                tokenizer(reference_answer, add_special_tokens=False).input_ids
            )
            expected_output_len = ref_out_len if dynamic_output else output_len

            # Validate sequence lengths.
            if not is_valid_sequence(prompt_len, expected_output_len):
                continue

            sampled_requests.append(
                SampleRequest(
                    prompt=prompt_formatted,
                    prompt_len=prompt_len,
                    expected_output_len=expected_output_len,
3379
                    request_id=request_id_prefix + str(ind),
3380
3381
                )
            )
3382
            ind += 1
3383

3384
3385
3386
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3387
        return sampled_requests
3388
3389
3390
3391
3392
3393
3394
3395


# -----------------------------------------------------------------------------
# Prefix Repetition Dataset Implementation
# -----------------------------------------------------------------------------


class PrefixRepetitionRandomDataset(BenchmarkDataset):
3396
    # Default values copied from benchmark_serving.py for the repeated prefix
3397
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
3408
3409
3410
3411
3412
    # dataset.
    DEFAULT_PREFIX_LEN = 256
    DEFAULT_SUFFIX_LEN = 256
    DEFAULT_NUM_PREFIXES = 10
    DEFAULT_OUTPUT_LEN = 128

    def __init__(
        self,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        random.seed(self.random_seed)
        np.random.seed(self.random_seed)

    def sample(
        self,
3413
        tokenizer: TokenizerLike,
3414
        num_requests: int,
3415
3416
        request_id_prefix: str = "",
        no_oversample: bool = False,
3417
3418
3419
3420
3421
3422
3423
3424
3425
3426
3427
3428
3429
3430
        prefix_len: int = DEFAULT_PREFIX_LEN,
        suffix_len: int = DEFAULT_SUFFIX_LEN,
        num_prefixes: int = DEFAULT_NUM_PREFIXES,
        output_len: int = DEFAULT_OUTPUT_LEN,
        **kwargs,
    ) -> list[SampleRequest]:
        vocab_size = tokenizer.vocab_size
        prompts_per_prefix = num_requests // num_prefixes
        if prompts_per_prefix == 0:
            raise ValueError(
                f"num_requests ({num_requests}) must be greater than or equal "
                f"to num_prefixes ({num_prefixes})"
            )

3431
        def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]:
3432
3433
3434
            """Generate tokens that decode and re-encode to exactly
            target_length."""
            # Generate random tokens
3435
            tokens = np.random.randint(0, vocab_size, size=target_length).tolist()
3436

3437
            _, adjusted_tokens, token_mismatch = gen_prompt_decode_to_target_len(  # noqa: E501
3438
3439
3440
3441
3442
3443
                tokenizer=tokenizer,
                token_sequence=tokens,
                target_token_len=target_length,
                add_special_tokens=False,
            )
            return adjusted_tokens, token_mismatch
3444
3445

        requests = []
3446
        token_mismatch_total = 0
3447
        for _ in range(num_prefixes):
3448
3449
            prefix_tokens, prefix_mismatch = _generate_exact_length_tokens(prefix_len)
            token_mismatch_total += prefix_mismatch
3450
3451

            for _ in range(prompts_per_prefix):
3452
                suffix_tokens, suffix_mismatch = _generate_exact_length_tokens(
3453
                    suffix_len
3454
                )
3455
                token_mismatch_total += suffix_mismatch
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
                combined_tokens = prefix_tokens + suffix_tokens
                prompt = tokenizer.decode(combined_tokens)
                prompt_len = len(combined_tokens)
                requests.append(
                    SampleRequest(
                        prompt=prompt,
                        prompt_len=prompt_len,
                        expected_output_len=output_len,
                    )
                )

3467
3468
3469
3470
3471
3472
3473
3474
3475
3476
        if token_mismatch_total != 0:
            sign = "more" if token_mismatch_total > 0 else "fewer"
            logger.warning(
                "Across all generated prompts, there were %d %s tokens "
                "than expected after decoding and re-encoding. This is "
                "expected due to the imperfect nature of the sampling "
                "procedure.",
                abs(token_mismatch_total),
                sign,
            )
3477
3478
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(requests)
3479
        return requests
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
3490
3491


# -----------------------------------------------------------------------------
# MMStar Dataset Implementation
# -----------------------------------------------------------------------------


class MMStarDataset(HuggingFaceDataset):
    """
    Lin-Chen/MMStar: https://huggingface.co/datasets/Lin-Chen/MMStar
    refer to: https://github.com/sgl-project/SpecForge/pull/106
    """
3492

3493
3494
3495
3496
3497
3498
    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {"Lin-Chen/MMStar"}
    IS_MULTIMODAL = True

    def sample(
        self,
3499
        tokenizer: TokenizerLike,
3500
3501
3502
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
3503
3504
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
3505
3506
3507
        **kwargs,
    ) -> list[SampleRequest]:
        # If --hf-output-len is not set, use the default output length.
3508
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
3509
3510
3511
3512
3513
3514
3515
3516
3517
3518
3519
3520
3521
3522
3523
3524
3525
        sampled_requests: list[SampleRequest] = []

        for ind, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
                break
            # Split the question text from options
            # (keep only the part before "Options:").
            full_q: str = item.get("question", "")
            question_text = full_q.split("Options:", 1)[0].strip()

            # Multimodal image content.
            mm_content = process_image(item["image"])

            # Compute prompt token length (note: this is plain text length
            # if enable_multimodal_chat is False).
            prompt_len = len(tokenizer(question_text).input_ids)

3526
            prompt: str | list[dict]
3527
3528
3529
3530
3531
3532
3533
3534
3535
3536
3537
3538
3539
3540
3541
3542
3543
3544
3545
3546
3547
3548
3549
3550
3551
3552
3553
            if enable_multimodal_chat:
                # If multimodal content should be embedded in the chat message,
                # convert to [{"role":"user","content":[...]}]
                prompt = self.apply_multimodal_chat_transformation(
                    question_text, mm_content
                )
                mm_for_request = None  # Already embedded in chat content.
            else:
                # Default: prompt is plain text,
                # image is in mm_content for the bench to assemble.
                prompt = question_text
                mm_for_request = mm_content

            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_for_request,
                    request_id=request_id_prefix + str(ind),
                )
            )

        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
        return sampled_requests