vllm-project--vllm
wehub-resource-sync
102 行
3.5 KiB
Python
102 行
3.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Shared utilities for benchmark dataset sampling.
|
|
"""
|
|
|
|
import logging
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
from vllm.tokenizers import TokenizerLike
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Type alias: a single float applies to both ISL and OSL; a dict allows
|
|
# specifying them independently via ``{"input": …, "output": …}``.
|
|
RangeRatio = float | dict[str, float]
|
|
|
|
|
|
def _resolve_range_ratios(
|
|
range_ratio: RangeRatio,
|
|
) -> tuple[float, float]:
|
|
"""Return ``(input_range_ratio, output_range_ratio)`` from *range_ratio*.
|
|
|
|
*range_ratio* is either a single float (used for both input and output)
|
|
or a dict with ``"input"`` and ``"output"`` keys.
|
|
"""
|
|
if isinstance(range_ratio, dict):
|
|
try:
|
|
return float(range_ratio["input"]), float(range_ratio["output"])
|
|
except KeyError as exc:
|
|
raise ValueError(
|
|
"When range_ratio is a dict it must contain 'input' and "
|
|
f"'output' keys, got: {sorted(range_ratio)}"
|
|
) from exc
|
|
ratio = float(range_ratio)
|
|
return ratio, ratio
|
|
|
|
|
|
def get_sampling_params(
|
|
rng: np.random.Generator,
|
|
num_requests: int,
|
|
range_ratio: RangeRatio,
|
|
input_len: int,
|
|
output_len: int,
|
|
tokenizer: TokenizerLike,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""
|
|
Sample per-request input/output token lengths and vocab offsets.
|
|
|
|
Lengths are drawn uniformly from integer ranges around the configured
|
|
means, controlled by *range_ratio*. It may be a single ``float``
|
|
(applied to both input and output) or a ``dict`` with ``"input"`` and
|
|
``"output"`` keys for independent control.
|
|
|
|
Tokenizer special tokens are subtracted from ``input_len`` before
|
|
computing the sampling interval.
|
|
|
|
Returns:
|
|
(input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of
|
|
shape ``(num_requests,)``.
|
|
"""
|
|
input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio)
|
|
|
|
if not (0.0 <= input_range_ratio < 1.0):
|
|
raise ValueError("input_range_ratio must be in [0, 1).")
|
|
if not (0.0 <= output_range_ratio < 1.0):
|
|
raise ValueError("output_range_ratio must be in [0, 1).")
|
|
num_special_tokens = int(tokenizer.num_special_tokens_to_add())
|
|
real_input_len = max(0, int(input_len) - num_special_tokens)
|
|
input_low = math.floor(real_input_len * (1 - input_range_ratio))
|
|
input_high = math.ceil(real_input_len * (1 + input_range_ratio))
|
|
output_low = math.floor(output_len * (1 - output_range_ratio))
|
|
output_high = math.ceil(output_len * (1 + output_range_ratio))
|
|
# Ensure the lower bound for output length is at least 1 to
|
|
# prevent sampling 0 tokens.
|
|
output_low = max(output_low, 1)
|
|
output_high = max(output_high, 1)
|
|
|
|
if input_low > input_high:
|
|
raise ValueError(
|
|
f"Invalid input sampling interval: low={input_low} > high={input_high}"
|
|
)
|
|
if output_low > output_high:
|
|
raise ValueError(
|
|
f"Invalid output sampling interval: low={output_low} > high={output_high}"
|
|
)
|
|
|
|
logger.info(
|
|
"Sampling input_len from [%s, %s] and output_len from [%s, %s]",
|
|
input_low,
|
|
input_high,
|
|
output_low,
|
|
output_high,
|
|
)
|
|
|
|
input_lens = rng.integers(input_low, input_high + 1, size=num_requests)
|
|
output_lens = rng.integers(output_low, output_high + 1, size=num_requests)
|
|
offsets = rng.integers(0, tokenizer.vocab_size, size=num_requests)
|
|
return input_lens, output_lens, offsets
|