datasets.py 132 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 pathlib import Path
29
from tempfile import NamedTemporaryFile
30
from typing import Any, cast
31
32

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

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

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

63
64
65

logger = logging.getLogger(__name__)

66
67
DEFAULT_NUM_PROMPTS = 1000

68
69
70
71
72
73
74

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

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


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


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

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

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

    def apply_multimodal_chat_transformation(
117
118
        self,
        prompt: str,
119
        mm_content: MultiModalDataDict | dict | list[dict] | None = None,
120
    ) -> list[dict]:
121
122
123
124
125
126
127
        """
        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:
128
129
130
131
132
            if isinstance(mm_content, list):
                content.extend(cast(list[dict[str, Any]], mm_content))
            elif isinstance(mm_content, dict):
                content.append(mm_content)
            else:
133
                raise TypeError(
134
                    f"Could not process multimodal content of type: {type(mm_content)}"
135
                )
136
137
138
139
140
141
142
143
144
145
146
147
148
        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
149
        raise NotImplementedError("load_data must be implemented in subclasses.")
150
151
152

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

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

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

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

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

246
    @abstractmethod
247
248
    def sample(
        self,
249
        tokenizer: TokenizerLike,
250
251
252
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
253
        **kwargs,
254
    ) -> list[SampleRequest]:
255
256
257
258
259
260
261
        """
        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:
262
            tokenizer (TokenizerLike): The tokenizer to be used
263
                for processing the dataset's text.
264
            num_requests (int): The number of sample requests to generate.
265
            request_id_prefix (str): The prefix of request_id.
266
267
268
269
270
271
272

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

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

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

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

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

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

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

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


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


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


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

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

    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.

372
373
374
375
376
    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.
377
378
379
380

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

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

401
    raise ValueError(
402
403
        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."
404
    )
405
406


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

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

443
444

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

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

463
464
465
    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.
466
467
468
469
470
    """
    remain_num_try = max_retry
    token_mismatch = 0
    while True:
        prompt = tokenizer.decode(token_sequence)
471
        token_sequence = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
472
473
474
475
        if remain_num_try <= 0:
            if len(token_sequence) != target_token_len:
                token_mismatch = len(token_sequence) - target_token_len
            break
476

477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
        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

500

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

505

506
class RandomDataset(BenchmarkDataset):
507
508
509
510
511
512
513
514
515
516
517
518
    """
    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.
    """
519

520
521
522
523
524
525
    # 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

526
    def __init__(self, **kwargs) -> None:
527
        super().__init__(**kwargs)
528
529
530
531
        # 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)
532
533
534

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

551
552
        num_special = int(tokenizer.num_special_tokens_to_add())
        real_input_len = max(0, int(input_len) - num_special)
553
554
555
        min_sampled_input = math.floor(
            real_input_len * (1.0 - float(resolved_input_rr))
        )
556
557
558
559
        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 "
560
561
                f"tokens {num_special} and "
                f"input range ratio {resolved_input_rr}, "
562
563
                "the minimum possible total input tokens (prefix + sampled) is "
                f"{min_total_input}. Increase --random-input-len and/or "
564
565
566
567
568
569
570
571
572
573
574
575
                "--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,
576
577
578
        )

        vocab_size = tokenizer.vocab_size
579
580
581
582
583
        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
584
        prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)
585

586
        requests = []
587
        token_mismatch_total = 0
588
        for i in range(num_requests):
589
            prompt, total_input_len, token_mismatch = self.generate_token_sequence(  # noqa: E501
590
591
592
593
594
595
596
                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,
597
                allowed_tokens=allowed_tokens,
598
            )
599
            token_mismatch_total += token_mismatch
600
601
602
603
604
605
            lora_req = self.get_lora_request(
                index=i,
                max_loras=max_loras,
                lora_path=lora_path,
                lora_assignment=lora_assignment,
            )
606
607
608
609
610
            requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=total_input_len,
                    expected_output_len=int(output_lens[i]),
611
                    lora_request=lora_req,
612
613
614
                    request_id=request_id_prefix + str(i),
                )
            )
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
        # 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
630

631
632
633
634
635
636
637
638
639
640
641
        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,
            )

642
643
644
        return requests

    def get_prefix(
645
        self,
646
        tokenizer: TokenizerLike,
647
648
        allowed_tokens: np.ndarray,
        prefix_len: int,
649
650
651
652
    ) -> list[int]:
        """
        Get the prefix for the dataset.
        """
653
654
655
656
657
658
659
660
661
662
663
664
        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,
665
        )
666
667
668
669
670
671
672
673
674
675
        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
676

677
678
679
    def generate_token_sequence(
        self,
        *,
680
        tokenizer: TokenizerLike,
681
682
683
684
685
686
        prefix_token_ids: list[int],
        prefix_len: int,
        vocab_size: int,
        input_len: int,
        offset: int,
        index: int,
687
        allowed_tokens: np.ndarray,
688
    ) -> tuple[str, int, int]:
689
690
691
692
693
694
695
696
697
698
        """
        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,
699
        the encoded sequence is truncated before being decoded again.
700
        """
701
702
703
704
705
        # 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()
706
707
708
709
        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)
710
        prompt, adjusted_token_sequence, token_mismatch = (
711
            gen_prompt_decode_to_target_len(
712
713
714
715
716
717
                tokenizer=tokenizer,
                token_sequence=token_sequence,
                target_token_len=total_input_len,
                add_special_tokens=False,
                rng=self._rng,
            )
718
719
720
        )
        total_input_len = len(adjusted_token_sequence)
        return prompt, total_input_len, token_mismatch
721
722


723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
# -----------------------------------------------------------------------------
# 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,
740
        tokenizer: TokenizerLike,
741
742
        num_requests: int,
        request_id_prefix: str = "",
743
744
745
        no_oversample: bool = False,
        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
746
        input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
747
        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
748
749
750
751
752
753
754
755
        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

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

        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

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

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

        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,
798
                allowed_tokens=allowed_tokens,
799
800
801
802
803
804
805
806
807
808
809
810
811
            )
        )

        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,
812
                allowed_tokens=allowed_tokens,
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
847
            )
            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


848
849
850
851
# -----------------------------------------------------------------------------
# MultiModalDataset Implementation
# -----------------------------------------------------------------------------

852

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

    Status:
    - Images: supported via synthetic RGB data.
859
    - Video: supported via synthetic RGB data.
860
861
862
863
864
865
866
867
    - 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
868
       mapping (height, width, num_frames) → probability. We treat
869
       `num_frames`=1 as image and `num_frames` > 1 as video.
870
       Entries with zero probability are removed and the rest are renormalized
871
872
873
874
875
876
877
       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}
878
879
      - Two image buckets (`num_frames`=1) and one video bucket
      (`num_frames`=16).
880
881
882
883
    OBS.: Only image sampling is supported for now.
    """

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

    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.
900
901
902

        NOTE: iid pixel sampling results in worst-case compression
        (good for stressing I/O), but very unlike real photos.
903
904
905
906
907
908
909
910
911
912
913
        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)

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

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

924
925
926
927
928
929
930
931
932
933
934
        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

935
        with NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
            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}
957
958
959
960
961
962
963
964
965
966

    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}")

967
968
969
    def normalize_bucket_config(
        self, bucket_config: dict[tuple[int, int, int], float]
    ) -> dict[tuple[int, int, int], float]:
970
971
972
973
974
975
976
977
978
979
980
        """
        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:
981
982
983
            raise ValueError(
                "Got invalid bucket config. Bucket config values must be non-zero."
            )
984
985
986
987
        # 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()}

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

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

    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:
1031
1032
1033
1034
1035
                raise ValueError(
                    f"Modality {modality} is not in "
                    f"limit_mm_per_prompt: "
                    f"{limit_mm_per_prompt.keys()}"
                )
1036

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

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

        # 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(
1060
            sum(limit_mm_per_prompt.values()),
1061
            math.ceil(base_items_per_request * (1 + num_mm_items_range_ratio)),
1062
1063
1064
        )
        # Ensure min num mm items is at least 0
        min_num_mm_items = max(
1065
            0, math.floor(base_items_per_request * (1 - num_mm_items_range_ratio))
1066
1067
1068
        )
        # Raise error if min num mm items is greater than max num mm items
        if min_num_mm_items > max_num_mm_items:
1069
1070
1071
1072
            raise ValueError(
                f"Min num mm items is greater than max mm items: "
                f"{min_num_mm_items} > {max_num_mm_items}"
            )
1073

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

        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],
1093
    ) -> Iterator[tuple[int, int, int]]:
1094
1095
1096
1097
1098
        """
        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.
1099
1100
        Loop until the number of multimodal items sampled is equal to
        request_num_mm_items or limit of multimodal items per prompt
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
        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)
1112
        )
1113
1114
1115
1116
        # If request_num_mm_items is 0, yield an empty iterator
        if request_num_mm_items == 0:
            return
        # Initialize modality counters
1117
        modality_counter = {self.map_config_to_modality(k): 0 for k in bucket_config}
1118
1119
1120
1121
1122
        # 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
1123
1124
1125
            mm_item_config = self._rng.choice(
                list(bucket_config_copy.keys()), p=list(bucket_config_copy.values())
            )
1126
1127
1128
1129
            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
1130
                yield (mm_item_config)
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
            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()):
1141
1142
1143
                    logger.warning(
                        "Exhausted all multimodal items of modality %s", modality
                    )
1144
1145
                    break
                # Renormalize the bucket config
1146
                bucket_config_copy = self.normalize_bucket_config(bucket_config_copy)
1147
1148
1149

    def sample(
        self,
1150
        tokenizer: TokenizerLike,
1151
1152
        num_requests: int,
        request_id_prefix: str = "",
1153
        no_oversample: bool = False,
1154
        prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN,
1155
        range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO,
1156
1157
        input_len: int = RandomDataset.DEFAULT_INPUT_LEN,
        output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN,
1158
        batchsize: int = 1,
1159
1160
1161
        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,
1162
1163
1164
        bucket_config: dict[
            tuple[int, int, int], float
        ] = DEFAULT_MM_ITEM_BUCKET_CONFIG,
1165
1166
1167
        enable_multimodal_chat: bool = DEFAULT_ENABLE_MULTIMODAL_CHAT,
        **kwargs,
    ) -> list[SampleRequest]:
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
        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,
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
        )

        (
            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
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
        # 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
1211
        prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len)
1212
1213
        # Add synthetic multimodal items to each request
        mm_requests = []
1214
        token_mismatch_total = 0
1215
        for i in range(num_requests):
1216
            prompt, total_input_len, token_mismatch = self.generate_token_sequence(  # noqa: E501
1217
1218
1219
1220
1221
1222
1223
                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,
1224
                allowed_tokens=allowed_tokens,
1225
            )
1226
            token_mismatch_total += token_mismatch
1227
1228
1229
1230
1231
1232
1233
1234
            # 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,
            )

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

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

        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,
            )

1278
        return mm_requests
1279

1280

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

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

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


1376
def add_dataset_parser(parser: FlexibleArgumentParser):
1377
1378
1379
1380
1381
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="Trust remote code from huggingface",
    )
1382
1383
1384
1385
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--num-prompts",
        type=int,
1386
        default=DEFAULT_NUM_PROMPTS,
1387
1388
1389
1390
1391
1392
        help="Number of prompts to process.",
    )
    parser.add_argument(
        "--dataset-name",
        type=str,
        default="random",
1393
        choices=[
1394
1395
1396
1397
1398
            "sharegpt",
            "burstgpt",
            "sonnet",
            "random",
            "random-mm",
1399
            "random-rerank",
1400
1401
            "hf",
            "custom",
1402
            "custom_mm",
1403
1404
            "prefix_repetition",
            "spec_bench",
1405
            "speed_bench",
1406
        ],
1407
1408
        help="Name of the dataset to benchmark on.",
    )
1409
1410
1411
1412
1413
    parser.add_argument(
        "--no-stream",
        action="store_true",
        help="Do not load the dataset in streaming mode.",
    )
1414
1415
1416
1417
    parser.add_argument(
        "--dataset-path",
        type=str,
        default=None,
1418
1419
        help="Path to the sharegpt/sonnet dataset or the HF dataset ID if "
        "using HF dataset.",
1420
    )
1421
1422
1423
    parser.add_argument(
        "--no-oversample",
        action="store_true",
1424
        help="Do not oversample if the dataset has fewer samples than num-prompts.",
1425
    )
1426
1427
1428
    parser.add_argument(
        "--skip-chat-template",
        action="store_true",
1429
        help="Skip applying chat template to prompt for datasets that support it.",
1430
    )
1431
1432
1433
1434
1435
    parser.add_argument(
        "--enable-multimodal-chat",
        action="store_true",
        help="Enable multimodal chat transformation for datasets that support it.",
    )
1436
1437
1438
1439
1440
    parser.add_argument(
        "--disable-shuffle",
        action="store_true",
        help="Disable shuffling of dataset samples for deterministic ordering.",
    )
1441
1442
1443
1444
1445
1446
1447

    # 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,
1448
1449
1450
        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.",
1451
1452
    )

1453
1454
1455
1456
1457
    spec_bench_group = parser.add_argument_group("spec bench dataset options")
    spec_bench_group.add_argument(
        "--spec-bench-output-len",
        type=int,
        default=256,
1458
        help="Num of output tokens per request, used only for spec bench dataset.",
1459
1460
1461
1462
1463
    )
    spec_bench_group.add_argument(
        "--spec-bench-category",
        type=str,
        default=None,
1464
        help="Category for spec bench dataset. If None, use all categories.",
1465
1466
    )

1467
1468
1469
1470
1471
    sonnet_group = parser.add_argument_group("sonnet dataset options")
    sonnet_group.add_argument(
        "--sonnet-input-len",
        type=int,
        default=550,
1472
        help="Number of input tokens per request, used only for sonnet dataset.",
1473
1474
1475
1476
1477
    )
    sonnet_group.add_argument(
        "--sonnet-output-len",
        type=int,
        default=150,
1478
        help="Number of output tokens per request, used only for sonnet dataset.",
1479
1480
1481
1482
1483
    )
    sonnet_group.add_argument(
        "--sonnet-prefix-len",
        type=int,
        default=200,
1484
        help="Number of prefix tokens per request, used only for sonnet dataset.",
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
    )

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

1496
1497
1498
1499
1500
    blazedit_group = parser.add_argument_group("blazedit dataset options")
    blazedit_group.add_argument(
        "--blazedit-min-distance",
        type=float,
        default=0.0,
1501
        help="Minimum distance for blazedit dataset. Min: 0, Max: 1.0",
1502
1503
1504
1505
1506
    )
    blazedit_group.add_argument(
        "--blazedit-max-distance",
        type=float,
        default=1.0,
1507
        help="Maximum distance for blazedit dataset. Min: 0, Max: 1.0",
1508
1509
    )

1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
    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.",
    )

1524
    random_group = parser.add_argument_group("random dataset options")
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
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
    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.",
    )

1589
1590
1591
    speed_bench_group = parser.add_argument_group(
        "speed bench dataset options", description=SpeedBench.__doc__
    )
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
    speed_bench_group.add_argument(
        "--speed-bench-dataset-subset",
        type=str,
        default="qualitative",
        choices={
            "qualitative",
            "throughput_1k",
            "throughput_2k",
            "throughput_8k",
            "throughput_16k",
            "throughput_32k",
        },
        help="Subset of the SPEED-Bench dataset.",
    )
    speed_bench_group.add_argument(
        "--speed-bench-output-len",
        type=int,
        default=4096,
        help="Num of output tokens per request, used only for speed bench dataset.",
    )
    speed_bench_group.add_argument(
        "--speed-bench-category",
        type=str,
        default=None,
        help="Category for speed bench dataset. If None, use all categories.",
    )

1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633

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(
1634
1635
1636
        "--random-input-len",
        type=int,
        default=1024,
1637
        help="Number of input tokens per request, used only for random sampling.",
1638
    )
1639
    parser_or_group.add_argument(
1640
1641
1642
        "--random-output-len",
        type=int,
        default=128,
1643
        help="Number of output tokens per request, used only for random sampling.",
1644
    )
1645
    parser_or_group.add_argument(
1646
        "--random-range-ratio",
1647
1648
        type=str,
        default="0.0",
1649
        help="Range ratio for sampling input/output length, "
1650
1651
1652
        "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).",
1653
    )
1654
    parser_or_group.add_argument(
1655
1656
1657
        "--random-prefix-len",
        type=int,
        default=0,
1658
1659
1660
1661
1662
1663
1664
1665
        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)]."
        ),
1666
    )
1667
    parser_or_group.add_argument(
1668
1669
1670
        "--random-batch-size",
        type=int,
        default=1,
1671
        help=("Batch size for random sampling. Only used for embeddings benchmark."),
1672
    )
1673
    parser_or_group.add_argument(
1674
1675
1676
1677
1678
1679
1680
        "--no-reranker",
        action="store_true",
        help=(
            "Whether the model supports reranking natively."
            " Only used for reranker benchmark."
        ),
    )
1681

1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694

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(
1695
1696
1697
1698
1699
1700
1701
1702
1703
        "--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."
        ),
    )
1704
    parser_or_group.add_argument(
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
        "--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."
        ),
    )
1719
    parser_or_group.add_argument(
1720
1721
1722
1723
1724
        "--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. "
1725
            '\'{"image": 3, "video": 0}\'. The sampled per-request item '
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
            "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)
1742
1743
1744
1745
1746
                if not (
                    isinstance(key, tuple)
                    and len(key) == 3
                    and all(isinstance(x, int) for x in key)
                ):
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
                    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.")

1763
    parser_or_group.add_argument(
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
        "--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."
        ),
1782
1783
    )

1784

1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
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)


1797
def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
1798
1799
1800
    if not hasattr(args, "request_id_prefix"):
        args.request_id_prefix = ""

1801
1802
1803
    if hasattr(args, "random_range_ratio") and isinstance(args.random_range_ratio, str):
        args.random_range_ratio = _parse_range_ratio(args.random_range_ratio)

1804
    if args.dataset_name == "custom":
1805
1806
1807
        dataset = CustomDataset(
            dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
        )
1808
1809
1810
1811
        input_requests = dataset.sample(
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
            output_len=args.custom_output_len,
1812
            skip_chat_template=args.skip_chat_template,
1813
            request_id_prefix=args.request_id_prefix,
1814
            no_oversample=args.no_oversample,
1815
1816
        )

1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
    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,
        )

1830
    elif args.dataset_name == "sonnet":
1831
1832
1833
        dataset = SonnetDataset(
            dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle
        )
1834
        # For the "sonnet" dataset, formatting depends on the backend.
1835
        if args.backend == "openai-chat":
1836
1837
1838
1839
1840
1841
1842
            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,
1843
                request_id_prefix=args.request_id_prefix,
1844
                no_oversample=args.no_oversample,
1845
1846
1847
            )
        else:
            assert tokenizer.chat_template or tokenizer.default_chat_template, (
1848
1849
                "Tokenizer/model must have chat template for sonnet dataset."
            )
1850
1851
1852
1853
1854
1855
1856
            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,
1857
                request_id_prefix=args.request_id_prefix,
1858
                no_oversample=args.no_oversample,
1859
1860
1861
1862
1863
            )

    elif args.dataset_name == "hf":
        # all following datasets are implemented from the
        # HuggingFaceDataset base class
1864
        hf_kwargs = {}
1865
1866
1867
1868
        if (
            args.dataset_path in VisionArenaDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in VisionArenaDataset.SUPPORTED_DATASET_PATHS
        ):
1869
            dataset_class = VisionArenaDataset
1870
            args.hf_split = args.hf_split if args.hf_split else "train"
1871
            args.hf_subset = None
1872
1873
1874
1875
1876
        elif (
            args.dataset_path in MMVUDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MMVUDataset.SUPPORTED_DATASET_PATHS
        ):
            dataset_class = MMVUDataset
1877
            args.hf_split = args.hf_split if args.hf_split else "validation"
1878
            args.hf_subset = None
1879
1880
1881
1882
        elif (
            args.dataset_path in InstructCoderDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in InstructCoderDataset.SUPPORTED_DATASET_PATHS
        ):
1883
            dataset_class = InstructCoderDataset
1884
            args.hf_split = args.hf_split if args.hf_split else "train"
1885
1886
1887
1888
        elif (
            args.dataset_path in MTBenchDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MTBenchDataset.SUPPORTED_DATASET_PATHS
        ):
1889
            dataset_class = MTBenchDataset
1890
            args.hf_split = args.hf_split if args.hf_split else "train"
1891
1892
1893
1894
1895
        elif (
            args.dataset_path in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
        ):
            dataset_class = MultiModalConversationDataset
1896
1897
1898
1899
        elif (
            args.dataset_path in ConversationDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in ConversationDataset.SUPPORTED_DATASET_PATHS
        ):
1900
            dataset_class = ConversationDataset
1901
1902
1903
1904
        elif (
            args.dataset_path in AIMODataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in AIMODataset.SUPPORTED_DATASET_PATHS
        ):
1905
            dataset_class = AIMODataset
1906
            args.hf_split = args.hf_split if args.hf_split else "train"
1907
        elif (
1908
            args.dataset_path in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS  # noqa: E501
1909
1910
            or args.hf_name in NextEditPredictionDataset.SUPPORTED_DATASET_PATHS
        ):
1911
            dataset_class = NextEditPredictionDataset
1912
            args.hf_split = args.hf_split if args.hf_split else "train"
1913
1914
1915
1916
        elif (
            args.dataset_path in ASRDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in ASRDataset.SUPPORTED_DATASET_PATHS
        ):
1917
            dataset_class = ASRDataset
1918
1919
1920
1921
1922
            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,
            }
1923
1924
        elif args.dataset_path in BlazeditDataset.SUPPORTED_DATASET_PATHS:
            dataset_class = BlazeditDataset
1925
            args.hf_split = args.hf_split if args.hf_split else "train"
1926
1927
1928
1929
            hf_kwargs = {
                "min_distance": args.blazedit_min_distance,
                "max_distance": args.blazedit_max_distance,
            }
1930
1931
1932
1933
        elif (
            args.dataset_path in MLPerfDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MLPerfDataset.SUPPORTED_DATASET_PATHS
        ):
1934
            dataset_class = MLPerfDataset
1935
            args.hf_split = args.hf_split if args.hf_split else "train"
1936
1937
1938
1939
1940
        elif (
            args.dataset_path in MMStarDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MMStarDataset.SUPPORTED_DATASET_PATHS
        ):
            dataset_class = MMStarDataset
1941
            args.hf_split = args.hf_split if args.hf_split else "val"
1942
            args.hf_subset = None
1943
        else:
1944
1945
1946
1947
1948
1949
1950
            supported_datasets = set(
                [
                    dataset_name
                    for cls in HuggingFaceDataset.__subclasses__()
                    for dataset_name in cls.SUPPORTED_DATASET_PATHS
                ]
            )
1951
1952
1953
1954
1955
            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 "
1956
1957
                "like to add support for additional dataset formats."
            )
1958

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

    else:
        # For datasets that follow a similar structure, use a mapping.
        dataset_mapping = {
1992
            "spec_bench": lambda: SpecBench(
1993
1994
1995
                dataset_path=args.dataset_path,
                category=args.spec_bench_category,
                disable_shuffle=args.disable_shuffle,
1996
            ).sample(
1997
1998
1999
                num_requests=args.num_prompts,
                tokenizer=tokenizer,
                output_len=args.spec_bench_output_len,
2000
                enable_multimodal_chat=args.enable_multimodal_chat,
2001
                request_id_prefix=args.request_id_prefix,
2002
                no_oversample=args.no_oversample,
2003
            ),
2004
            "sharegpt": lambda: ShareGPTDataset(
2005
2006
2007
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2008
2009
2010
2011
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                output_len=args.sharegpt_output_len,
2012
                enable_multimodal_chat=args.enable_multimodal_chat,
2013
                request_id_prefix=args.request_id_prefix,
2014
                no_oversample=args.no_oversample,
2015
2016
            ),
            "burstgpt": lambda: BurstGPTDataset(
2017
2018
2019
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2020
2021
2022
2023
            ).sample(
                tokenizer=tokenizer,
                num_requests=args.num_prompts,
                request_id_prefix=args.request_id_prefix,
2024
                no_oversample=args.no_oversample,
2025
2026
            ),
            "random": lambda: RandomDataset(
2027
2028
2029
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2030
            ).sample(
2031
2032
2033
2034
2035
2036
                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,
2037
                request_id_prefix=args.request_id_prefix,
2038
                batchsize=args.random_batch_size,
2039
                no_oversample=args.no_oversample,
2040
            ),
2041
            "random-mm": lambda: RandomMultiModalDataset(
2042
2043
2044
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
            ).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,
2056
                enable_multimodal_chat=args.enable_multimodal_chat,
2057
                request_id_prefix=args.request_id_prefix,
2058
                no_oversample=args.no_oversample,
2059
            ),
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
            "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,
            ),
2073
            "prefix_repetition": lambda: PrefixRepetitionRandomDataset(
2074
2075
2076
                random_seed=args.seed,
                dataset_path=args.dataset_path,
                disable_shuffle=args.disable_shuffle,
2077
2078
2079
2080
2081
2082
2083
            ).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,
2084
                request_id_prefix=args.request_id_prefix,
2085
                no_oversample=args.no_oversample,
2086
            ),
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
            "speed_bench": lambda: SpeedBench(
                dataset_path=args.dataset_path,
                dataset_subset=args.speed_bench_dataset_subset,
                category=args.speed_bench_category,
                disable_shuffle=args.disable_shuffle,
            ).sample(
                num_requests=args.num_prompts,
                tokenizer=tokenizer,
                output_len=args.speed_bench_output_len,
                enable_multimodal_chat=args.enable_multimodal_chat,
                request_id_prefix=args.request_id_prefix,
                no_oversample=args.no_oversample,
            ),
2100
2101
2102
        }

        try:
2103
            # Enforce endpoint compatibility for multimodal datasets.
2104
            if args.dataset_name == "random-mm" and args.backend not in ["openai-chat"]:
2105
2106
2107
2108
                raise ValueError(
                    "Multi-modal content (images) is only supported on "
                    "'openai-chat' backend."
                )
2109
2110
2111
2112
2113
2114
2115
            input_requests = dataset_mapping[args.dataset_name]()
        except KeyError as err:
            raise ValueError(f"Unknown dataset: {args.dataset_name}") from err

    return input_requests


2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
# -----------------------------------------------------------------------------
# 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.,
    ```
2126
2127
2128
    {"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}
2129
    ```
2130
2131
    Note that 'output_tokens' column is optional and has to be provided only if
    'custom-output-len' argument is None or -1.
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
    """

    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
2147
        self.data: list[dict] = []
2148
2149
2150

        # Load the JSONL file
        if self.dataset_path.endswith(".jsonl"):
2151
            jsonl_data = pd.read_json(path_or_buf=self.dataset_path, lines=True)
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164

            # 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(
2165
2166
                "Only JSONL format is supported for CustomDataset."
            )
2167
2168

        random.seed(self.random_seed)
2169
2170
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(self.data)
2171
2172
2173

    def sample(
        self,
2174
        tokenizer: TokenizerLike,
2175
        num_requests: int,
2176
2177
        request_id_prefix: str = "",
        no_oversample: bool = False,
2178
2179
2180
        lora_path: str | None = None,
        max_loras: int | None = None,
        output_len: int | None = None,
2181
2182
2183
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
        **kwargs,
2184
    ) -> list[SampleRequest]:
2185
2186
2187
2188
        # load all data if needed
        self.num_available_samples = len(self.data)
        if num_requests <= 0:
            num_requests = self.num_available_samples
2189
2190
2191
2192
2193
            logger.info(
                "num_requests is set to 0 or negative, "
                "so using all available samples: %d",
                num_requests,
            )
2194

2195
        sampled_requests: list[SampleRequest] = []
2196
        for i, item in enumerate(self.data):
2197
2198
2199
2200
            if len(sampled_requests) >= num_requests:
                break
            prompt = item["prompt"]

2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
            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,
2230
                    )
2231

2232
                prompt_len = len(tokenizer(prompt).input_ids)
2233
2234
2235
2236
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
2237
                    expected_output_len=new_output_len,
2238
                    request_id=request_id_prefix + str(i),
2239
2240
2241
2242
2243
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2244
2245
2246
2247

        return sampled_requests


2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
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,
2279
    ) -> list[SampleRequest]:
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
        # 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


2327
2328
2329
2330
2331
2332
2333
2334
# -----------------------------------------------------------------------------
# Spec Bench Dataset Implementation
# -----------------------------------------------------------------------------


class SpecBench(CustomDataset):
    """
    Implements the SpecBench dataset: https://github.com/hemingkx/Spec-Bench
2335
    Download the dataset using:
2336
    wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl
2337
    """  # noqa: E501
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350

    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
2351
        jsonl_data = pd.read_json(path_or_buf=self.dataset_path, lines=True)
2352
2353
2354
2355
2356
2357
2358

        # 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
2359
            if (not self.category) or (self.category == row["category"]):
2360
2361
2362
2363
                prompt = row["turns"][0]
                self.data.append({"prompt": prompt})

        random.seed(self.random_seed)
2364
2365
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(self.data)
2366

2367
2368
2369
    def sample(
        **kwargs,
    ) -> list[SampleRequest]:
2370
        # leverage CustomDataset sample
2371
2372
2373
        return super().sample(
            **kwargs,
        )
2374
2375


2376
2377
2378
2379
# -----------------------------------------------------------------------------
# Sonnet Dataset Implementation
# -----------------------------------------------------------------------------

2380

2381
2382
2383
@deprecated(
    "SonnetDataset is deprecated and will be removed in a future version.",
)
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
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,
2410
        tokenizer: TokenizerLike,
2411
        num_requests: int,
2412
2413
        request_id_prefix: str = "",
        no_oversample: bool = False,
2414
2415
2416
2417
2418
        prefix_len: int = DEFAULT_PREFIX_LEN,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        return_prompt_formatted: bool = False,
        **kwargs,
2419
    ) -> list[SampleRequest]:
2420
2421
        # Calculate average token length for a poem line.
        tokenized_lines = [tokenizer(line).input_ids for line in self.data]
2422
        avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines)
2423
2424
2425
2426

        # 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}]
2427
2428
2429
        base_fmt = tokenizer.apply_chat_template(
            base_msg, add_generation_prompt=True, tokenize=False
        )
2430
2431
2432
2433
        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 "
2434
2435
                f"({base_offset})."
            )
2436
2437
2438
2439
2440
2441

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

2442
        samples: list[SampleRequest] = []
2443
        ind = 0
2444
        while len(samples) < num_requests:
2445
2446
2447
            extra_lines = random.choices(
                self.data, k=num_input_lines - num_prefix_lines
            )
2448
2449
2450
            prompt = f"{base_prompt}{''.join(prefix_lines + extra_lines)}"
            msg = [{"role": "user", "content": prompt}]
            prompt_formatted = tokenizer.apply_chat_template(
2451
2452
                msg, add_generation_prompt=True, tokenize=False
            )
2453
2454
2455
2456
            prompt_len = len(tokenizer(prompt_formatted).input_ids)
            if prompt_len <= input_len:
                samples.append(
                    SampleRequest(
2457
                        prompt=prompt_formatted if return_prompt_formatted else prompt,
2458
2459
                        prompt_len=prompt_len,
                        expected_output_len=output_len,
2460
2461
2462
                        request_id=request_id_prefix + str(ind),
                    )
                )
2463
                ind += 1
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
        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()

2483
2484
2485
    def load_data(
        self,
    ):
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
        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):
2499
            data = self.data.sample(n=num_requests, random_state=self.random_seed)
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
        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,
2511
        tokenizer: TokenizerLike,
2512
        num_requests: int,
2513
        request_id_prefix: str = "",
2514
        no_oversample: bool = False,
2515
        lora_assignment: str = "random",
2516
2517
        max_loras: int | None = None,
        lora_path: str | None = None,
2518
2519
2520
2521
2522
2523
2524
        **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])
2525
2526
2527
2528
2529
            lora_req = self.get_lora_request(
                index=i,
                max_loras=max_loras,
                lora_path=lora_path,
                lora_assignment=lora_assignment,
2530
            )
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
            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,
2542
                    request_id=request_id_prefix + str(i),
2543
2544
                )
            )
2545
2546
2547
2548
2549
2550
2551
2552
2553
        return samples


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

2554
    SUPPORTED_DATASET_PATHS: set[str] | dict[str, Callable] = set()
2555
2556
2557
2558
2559

    def __init__(
        self,
        dataset_path: str,
        dataset_split: str,
2560
        no_stream: bool = False,
2561
2562
        dataset_subset: str | None = None,
        hf_name: str | None = None,
2563
        trust_remote_code: bool = False,
2564
2565
2566
2567
2568
2569
        **kwargs,
    ) -> None:
        super().__init__(dataset_path=dataset_path, **kwargs)

        self.dataset_split = dataset_split
        self.dataset_subset = dataset_subset
2570
        self.load_stream = not no_stream
2571
        self.hf_name = hf_name or dataset_path
2572
        self.trust_remote_code = trust_remote_code
2573
2574
2575
2576
2577
2578
2579
2580
        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,
2581
            streaming=self.load_stream,
2582
            trust_remote_code=self.trust_remote_code,
2583
        )
2584
2585
        if not getattr(self, "disable_shuffle", False):
            self.data = self.data.shuffle(seed=self.random_seed)
2586
2587
2588
2589
2590
2591
2592
2593


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


class ConversationDataset(HuggingFaceDataset):
2594
    """Dataset for text-only conversation data."""
2595

2596
    SUPPORTED_DATASET_PATHS = {
2597
        "Aeala/ShareGPT_Vicuna_unfiltered",
2598
    }
2599
2600
2601
2602
    IS_MULTIMODAL = False

    def sample(
        self,
2603
        tokenizer: TokenizerLike,
2604
2605
2606
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
2607
2608
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2609
        **kwargs,
2610
    ) -> list[SampleRequest]:
2611
2612
        # Filter examples with at least 2 conversations
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
2613
        sampled_requests: list[SampleRequest] = []
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
        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",
    }
2659
    IS_MULTIMODAL = True
2660

2661
2662
    def sample(
        self,
2663
        tokenizer: TokenizerLike,
2664
2665
2666
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
2667
2668
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2669
        **kwargs,
2670
    ) -> list[SampleRequest]:
2671
        # Filter examples with at least 2 conversations
2672
        filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2)
2673
        sampled_requests: list[SampleRequest] = []
2674
        ind = 0
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
        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
2689
            if dynamic_output and not is_valid_sequence(prompt_len, completion_len):
2690
                continue
2691
            mm_content = process_image(item["image"]) if "image" in item else None
2692
2693
2694
2695
            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
2696
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
2697
2698
2699
2700
2701
2702
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
2703
                    request_id=request_id_prefix + str(ind),
2704
2705
                )
            )
2706
            ind += 1
2707
2708
2709
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
        return sampled_requests


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


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

    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {
2725
2726
        "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"],
2727
    }
2728
    IS_MULTIMODAL = True
2729
2730
2731

    def sample(
        self,
2732
        tokenizer: TokenizerLike,
2733
        num_requests: int,
2734
        request_id_prefix: str = "",
2735
        no_oversample: bool = False,
2736
2737
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2738
        **kwargs,
2739
    ) -> list[SampleRequest]:
2740
2741
2742
2743
        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
        if parser_fn is None:
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")

2744
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2745

2746
        sampled_requests = []
2747
        for i, item in enumerate(self.data):
2748
2749
            if len(sampled_requests) >= num_requests:
                break
2750

2751
2752
            prompt = parser_fn(item)
            mm_content = process_image(item["images"][0])
2753
            prompt_len = len(tokenizer.encode(prompt))
2754
2755
2756
2757
            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
2758
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
2759

2760
2761
2762
2763
2764
2765
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=mm_content,
2766
                    request_id=request_id_prefix + str(i),
2767
2768
                )
            )
2769

2770
2771
2772
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2773
2774
2775
        return sampled_requests


2776
2777
2778
2779
2780
2781
2782
2783
class MMVUDataset(HuggingFaceDataset):
    """
    MMVU Dataset.
    https://huggingface.co/datasets/yale-nlp/MMVU
    """

    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {
2784
2785
2786
2787
2788
        "yale-nlp/MMVU": lambda x: (
            x["question"]
            + " "
            + (" ".join(f"{k}.{v}" for k, v in x["choices"].items()))
        ),
2789
2790
    }

2791
2792
2793
2794
2795
2796
2797
2798
    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")

2799
2800
    def sample(
        self,
2801
        tokenizer: TokenizerLike,
2802
2803
2804
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
2805
2806
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
2807
        **kwargs,
2808
    ) -> list[SampleRequest]:
2809
2810
2811
2812
        parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name)
        if parser_fn is None:
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")

2813
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2814

2815
2816
2817
2818
        sampled_requests = []
        for i, item in enumerate(self.data):
            if len(sampled_requests) >= num_requests:
                break
2819

2820
            prompt = parser_fn(item)
2821
2822
2823
            mm_content = process_video(
                item["video"].replace(self._remote_path_root, self._local_path_root)
            )
2824
            prompt_len = len(tokenizer.encode(prompt))
2825
2826
2827
2828
            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
2829
                prompt = self.apply_multimodal_chat_transformation(prompt, mm_content)
2830

2831
2832
2833
2834
2835
2836
2837
            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),
2838
2839
                )
            )
2840

2841
2842
2843
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2844
2845
2846
        return sampled_requests


2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
# -----------------------------------------------------------------------------
# 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",
    }

2867
2868
    def sample(
        self,
2869
        tokenizer: TokenizerLike,
2870
        num_requests: int,
2871
2872
        request_id_prefix: str = "",
        no_oversample: bool = False,
2873
        output_len: int | None = None,
2874
2875
2876
        enable_multimodal_chat: bool = False,
        skip_chat_template: bool = False,
        **kwargs,
2877
    ) -> list[SampleRequest]:
2878
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2879
        sampled_requests: list[SampleRequest] = []
2880
        for i, prompt in enumerate(self.sample_prompts(n=num_requests)):
2881
            # apply template
2882
2883
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
2884
                    [{"role": "user", "content": prompt}],
2885
2886
2887
                    add_generation_prompt=True,
                    tokenize=False,
                )
2888

2889
2890
2891
2892
2893
2894
            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
2895
                    request_id=request_id_prefix + str(i),
2896
2897
2898
2899
2900
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2901
2902
        return sampled_requests

2903
2904
2905
2906
2907
2908
2909
2910
    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

2911

2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
# -----------------------------------------------------------------------------
# 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,
2934
        tokenizer: TokenizerLike,
2935
        num_requests: int,
2936
2937
        request_id_prefix: str = "",
        no_oversample: bool = False,
2938
        output_len: int | None = None,
2939
        enable_multimodal_chat: bool = False,
2940
        skip_chat_template: bool = False,
2941
        **kwargs,
2942
    ) -> list[SampleRequest]:
2943
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
2944
        sampled_requests: list[SampleRequest] = []
2945

2946
        for i, item in enumerate(self.data):
2947
2948
2949
2950
2951
            if len(sampled_requests) >= num_requests:
                break
            prompt = item["turns"][0]

            # apply template
2952
2953
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
2954
                    [{"role": "user", "content": prompt}],
2955
2956
2957
                    add_generation_prompt=True,
                    tokenize=False,
                )
2958
2959
2960
2961
2962
2963
2964

            prompt_len = len(tokenizer(prompt).input_ids)
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
2965
                    request_id=request_id_prefix + str(i),
2966
2967
2968
2969
2970
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
2971
2972
2973
        return sampled_requests


2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
# -----------------------------------------------------------------------------
# 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,
3000
        tokenizer: TokenizerLike,
3001
        num_requests: int,
3002
        output_len: int | None = None,
3003
        skip_chat_template: bool = False,
3004
        request_id_prefix: str = "",
3005
        no_oversample: bool = False,
3006
3007
3008
        min_distance: float = 0.0,
        max_distance: float = 1.0,
        **kwargs,
3009
    ) -> list[SampleRequest]:
3010
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
3011
3012
3013
3014
3015
3016
3017
3018
3019
3020
3021
3022
        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
3023
3024

            # template copied from
3025
            # https://github.com/ise-uiuc/blazedit/blob/7765137e656fd62de877422d2e4cf8de51228054/dataset/create_refined_dataset.py#L94-L105 # noqa: E501
3026
            prompt = f"""Given a code file, please apply the change requests and generate the new file.
3027
3028
3029
3030
3031
3032
3033
3034
3035

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

Change request:
{change_request}

3036
Please generate the new code file in the "New file" section below."""  # noqa: E501
3037
3038

            # apply template
3039
3040
            if not skip_chat_template:
                prompt = tokenizer.apply_chat_template(
3041
                    [{"role": "user", "content": prompt}],
3042
3043
3044
                    add_generation_prompt=True,
                    tokenize=False,
                )
3045
3046
3047
3048
3049
3050
3051
3052
3053

            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),
3054
3055
3056
3057
3058
                )
            )
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3059

3060
3061
3062
        return sampled_requests


3063
3064
3065
3066
3067
3068
3069
3070
3071
# -----------------------------------------------------------------------------
# AIMO Dataset Implementation
# -----------------------------------------------------------------------------


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

3073
    SUPPORTED_DATASET_PATHS = {
3074
3075
3076
        "AI-MO/aimo-validation-aime",
        "AI-MO/NuminaMath-1.5",
        "AI-MO/NuminaMath-CoT",
3077
3078
    }

3079
3080
    def sample(
        self,
3081
        tokenizer: TokenizerLike,
3082
3083
3084
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
3085
        output_len: int | None = None,
3086
        **kwargs,
3087
3088
    ) -> list[SampleRequest]:
        sampled_requests: list[SampleRequest] = []
3089
        ind = 0
3090
3091
3092
3093
3094
        dynamic_output = output_len is None

        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break
3095
            prompt, completion = item["problem"], item["solution"]
3096
3097
3098
3099
3100
3101
3102

            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
3103
3104
3105
            if dynamic_output and not is_valid_sequence(
                prompt_len, completion_len, max_prompt_len=2048, max_total_len=32000
            ):
3106
3107
3108
3109
3110
3111
3112
                continue
            sampled_requests.append(
                SampleRequest(
                    prompt=prompt,
                    prompt_len=prompt_len,
                    expected_output_len=output_len,
                    multi_modal_data=None,
3113
                    request_id=request_id_prefix + str(ind),
3114
3115
                )
            )
3116
            ind += 1
3117
3118
3119
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3120
        return sampled_requests
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140


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

3141
"""  # noqa: E501
3142
3143
3144


def _format_zeta_prompt(
3145
3146
    sample: dict, original_start_marker: str = "<|editable_region_start|>"
) -> dict:
3147
    """Format the zeta prompt for the Next Edit Prediction (NEP) dataset.
3148
3149
3150

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

3153
    Args:
3154
        sample: The dataset sample containing events,
3155
            inputs, and outputs.
3156
3157
        original_start_marker: The marker indicating the
            start of the editable region. Defaults to
3158
            "<|editable_region_start|>".
3159

3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
3187
3188
    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,
    }

3189
3190
    def sample(
        self,
3191
        tokenizer: TokenizerLike,
3192
3193
3194
3195
3196
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
        **kwargs,
    ):
3197
        formatting_prompt_func = self.MAPPING_PROMPT_FUNCS.get(self.hf_name)
3198
        if formatting_prompt_func is None:
3199
            raise ValueError(f"Unsupported dataset path: {self.hf_name}")
3200
        samples = []
3201
        for i, sample in enumerate(self.data):
3202
3203
3204
3205
3206
3207
            sample = formatting_prompt_func(sample)
            samples.append(
                SampleRequest(
                    prompt=sample["prompt"],
                    prompt_len=len(tokenizer(sample["prompt"]).input_ids),
                    expected_output_len=len(
3208
3209
                        tokenizer(sample["expected_output"]).input_ids
                    ),
3210
                    request_id=request_id_prefix + str(i),
3211
3212
                )
            )
3213
3214
            if len(samples) >= num_requests:
                break
3215
3216
3217
        self.maybe_oversample_requests(
            samples, num_requests, request_id_prefix, no_oversample
        )
3218
        return samples
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
3240
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
3251
3252
3253


# -----------------------------------------------------------------------------
# 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",
    }

3254
    DEFAULT_OUTPUT_LEN = 1024
3255
3256
3257
3258
    IS_MULTIMODAL = True

    def sample(
        self,
3259
        tokenizer: TokenizerLike,
3260
        num_requests: int,
3261
        request_id_prefix: str = "",
3262
        no_oversample: bool = False,
3263
        output_len: int | None = None,
3264
        **kwargs,
3265
    ) -> list[SampleRequest]:
3266
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
Ekagra Ranjan's avatar
Ekagra Ranjan committed
3267
        if "openai" in getattr(tokenizer, "name_or_path", ""):
3268
3269
3270
            prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
        else:
            prompt = ""
3271
        prompt_len = len(tokenizer(prompt).input_ids)
3272
        sampled_requests: list[SampleRequest] = []
3273
        ind = 0
3274
        skipped = 0
3275
3276
3277
        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 = []
3278
3279
3280
3281
3282
        for item in self.data:
            if len(sampled_requests) >= num_requests:
                break
            audio = item["audio"]
            y, sr = audio["array"], audio["sampling_rate"]
3283
            duration_s = get_audio_duration(y=y, sr=sr)
3284
            if duration_s < asr_min_audio_len_sec or duration_s > asr_max_audio_len_sec:
3285
3286
3287
                skipped += 1
                continue

3288
            durations.append(duration_s)
3289
3290
3291
3292
3293
3294
3295
            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,
3296
                    request_id=request_id_prefix + str(ind),
3297
3298
                )
            )
3299
            ind += 1
3300
3301
3302
3303
3304
3305
3306
        if skipped:
            logger.warning(
                "%d samples discarded from dataset due to"
                " their length being greater than"
                " what Whisper supports.",
                skipped,
            )
3307
3308
3309
3310
3311
3312
3313
3314
3315
3316
3317
3318
3319
3320

        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,
        )

3321
3322
3323
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3324
        return sampled_requests
3325
3326
3327
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
3339
3340
3341
3342
3343
3344
3345
3346
3347
3348
3349
3350
3351
3352
3353
3354
3355
3356


# -----------------------------------------------------------------------------
# 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,
3357
        tokenizer: TokenizerLike,
3358
        num_requests: int,
3359
        request_id_prefix: str = "",
3360
        no_oversample: bool = False,
3361
        output_len: int | None = None,
3362
3363
3364
3365
3366
        **kwargs,
    ) -> list[SampleRequest]:
        # Force dynamic output length based on reference completion.
        dynamic_output = output_len is None
        sampled_requests: list[SampleRequest] = []
3367
        ind = 0
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378
3379
3380
3381
3382
3383
3384
3385
3386
3387
3388
3389
3390
3391
3392
3393
3394
3395
3396
3397
3398
3399
3400
3401

        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,
3402
                    request_id=request_id_prefix + str(ind),
3403
3404
                )
            )
3405
            ind += 1
3406

3407
3408
3409
        self.maybe_oversample_requests(
            sampled_requests, num_requests, request_id_prefix, no_oversample
        )
3410
        return sampled_requests
3411
3412
3413
3414
3415
3416
3417
3418


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


class PrefixRepetitionRandomDataset(BenchmarkDataset):
3419
    # Default values copied from benchmark_serving.py for the repeated prefix
3420
3421
3422
3423
3424
3425
3426
3427
3428
3429
3430
3431
3432
3433
3434
3435
    # 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,
3436
        tokenizer: TokenizerLike,
3437
        num_requests: int,
3438
3439
        request_id_prefix: str = "",
        no_oversample: bool = False,
3440
3441
3442
3443
3444
3445
3446
3447
3448
3449
3450
3451
3452
3453
        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})"
            )

3454
        def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]:
3455
3456
3457
            """Generate tokens that decode and re-encode to exactly
            target_length."""
            # Generate random tokens
3458
            tokens = np.random.randint(0, vocab_size, size=target_length).tolist()
3459

3460
            _, adjusted_tokens, token_mismatch = gen_prompt_decode_to_target_len(  # noqa: E501
3461
3462
3463
3464
3465
3466
                tokenizer=tokenizer,
                token_sequence=tokens,
                target_token_len=target_length,
                add_special_tokens=False,
            )
            return adjusted_tokens, token_mismatch
3467
3468

        requests = []
3469
        token_mismatch_total = 0
3470
        for _ in range(num_prefixes):
3471
3472
            prefix_tokens, prefix_mismatch = _generate_exact_length_tokens(prefix_len)
            token_mismatch_total += prefix_mismatch
3473
3474

            for _ in range(prompts_per_prefix):
3475
                suffix_tokens, suffix_mismatch = _generate_exact_length_tokens(
3476
                    suffix_len
3477
                )
3478
                token_mismatch_total += suffix_mismatch
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
                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,
                    )
                )

3490
3491
3492
3493
3494
3495
3496
3497
3498
3499
        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,
            )
3500
3501
        if not getattr(self, "disable_shuffle", False):
            random.shuffle(requests)
3502
        return requests
3503
3504
3505
3506
3507
3508
3509
3510
3511
3512
3513
3514


# -----------------------------------------------------------------------------
# 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
    """
3515

3516
3517
3518
3519
3520
3521
    DEFAULT_OUTPUT_LEN = 128
    SUPPORTED_DATASET_PATHS = {"Lin-Chen/MMStar"}
    IS_MULTIMODAL = True

    def sample(
        self,
3522
        tokenizer: TokenizerLike,
3523
3524
3525
        num_requests: int,
        request_id_prefix: str = "",
        no_oversample: bool = False,
3526
3527
        output_len: int | None = None,
        enable_multimodal_chat: bool = False,
3528
3529
3530
        **kwargs,
    ) -> list[SampleRequest]:
        # If --hf-output-len is not set, use the default output length.
3531
        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
3532
3533
3534
3535
3536
3537
3538
3539
3540
3541
3542
3543
3544
3545
3546
3547
3548
        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)

3549
            prompt: str | list[dict]
3550
3551
3552
3553
3554
3555
3556
3557
3558
3559
3560
3561
3562
3563
3564
3565
3566
3567
3568
3569
3570
3571
3572
3573
3574
3575
3576
            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
3577
3578
3579
3580
3581
3582
3583
3584
3585


# -----------------------------------------------------------------------------
# Speed Bench Dataset Implementation
# -----------------------------------------------------------------------------


class SpeedBench(CustomDataset):
    """
3586
3587
    SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench

3588
    Download the dataset using:
3589
3590

    `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -`
3591
3592
    """  # noqa: E501

3593
3594
    DOWNLOAD_SCRIPT_URL = "https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py"

3595
3596
3597
3598
3599
3600
3601
3602
3603
3604
    def __init__(self, **kwargs) -> None:
        self.dataset_subset = kwargs.pop("dataset_subset", "qualitative")
        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.")

3605
3606
3607
3608
3609
3610
3611
        if not Path(self.dataset_path).is_dir():
            raise ValueError(
                f"dataset_path {self.dataset_path} is not a directory. "
                f"Please make sure to download the dataset from HuggingFace using "
                f"`curl -LsSf {self.DOWNLOAD_SCRIPT_URL} | python3 -`"
            )

3612
3613
3614
3615
3616
3617
3618
3619
3620
3621
        self.data = []

        # Load the JSONL file
        jsonl_data = pd.read_json(
            path_or_buf=Path(self.dataset_path) / f"{self.dataset_subset}.jsonl",
            lines=True,
        )

        # check if the JSONL file has a 'turns' column
        if "messages" not in jsonl_data.columns:
3622
3623
3624
3625
3626
            raise ValueError(
                "JSONL file must contain a 'messages' column. "
                "Please make sure to download the dataset from HuggingFace using "
                f"`curl -LsSf {self.DOWNLOAD_SCRIPT_URL} | python3 -`"
            )
3627
3628
3629
3630
3631
3632
3633
3634
3635
3636

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

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