wehub-resource-sync
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from transformers import AutoTokenizer, PreTrainedTokenizerBase
|
|
|
|
from vllm.benchmarks.datasets import get_samples
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def hf_tokenizer() -> PreTrainedTokenizerBase:
|
|
return AutoTokenizer.from_pretrained("openai-community/gpt2")
|
|
|
|
|
|
def _write_jsonl(path: Path, n_rows: int) -> None:
|
|
with path.open("w") as f:
|
|
for i in range(n_rows):
|
|
f.write(json.dumps({"prompt": f"row {i}: unique prompt content."}) + "\n")
|
|
|
|
|
|
def _args_for_custom(dataset_path: str, seed: int) -> argparse.Namespace:
|
|
return argparse.Namespace(
|
|
dataset_name="custom",
|
|
dataset_path=dataset_path,
|
|
disable_shuffle=False,
|
|
num_prompts=30,
|
|
custom_output_len=32,
|
|
skip_chat_template=True,
|
|
no_oversample=False,
|
|
seed=seed,
|
|
request_id_prefix="",
|
|
)
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_custom_dataset_seed_propagates(
|
|
hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
|
|
) -> None:
|
|
"""--seed must control the CustomDataset shuffle used by get_samples.
|
|
|
|
Without the fix, CustomDataset was instantiated without random_seed,
|
|
so its load-time shuffle always used DEFAULT_SEED=0 regardless of
|
|
args.seed, causing every run with --dataset-name custom to pick the
|
|
same subset of rows from a larger file.
|
|
"""
|
|
jsonl = tmp_path / "data.jsonl"
|
|
_write_jsonl(jsonl, n_rows=60)
|
|
|
|
samples_a = get_samples(_args_for_custom(str(jsonl), seed=0), hf_tokenizer)
|
|
samples_b = get_samples(_args_for_custom(str(jsonl), seed=42), hf_tokenizer)
|
|
|
|
prompts_a = {s.prompt for s in samples_a}
|
|
prompts_b = {s.prompt for s in samples_b}
|
|
|
|
assert len(prompts_a) == 30
|
|
assert len(prompts_b) == 30
|
|
assert prompts_a != prompts_b
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_custom_dataset_same_seed_is_deterministic(
|
|
hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
|
|
) -> None:
|
|
"""Same --seed must yield the same CustomDataset subset."""
|
|
jsonl = tmp_path / "data.jsonl"
|
|
_write_jsonl(jsonl, n_rows=60)
|
|
|
|
samples_a = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer)
|
|
samples_b = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer)
|
|
|
|
prompts_a = [s.prompt for s in samples_a]
|
|
prompts_b = [s.prompt for s in samples_b]
|
|
|
|
assert prompts_a == prompts_b
|