wehub-resource-sync
102 lines
3.5 KiB
Python
102 lines
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
|