forked from Karylab-cklius/vllm
Harry MellorandGitHub
f1445f6dbd
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from transformers import AutoTokenizer, PreTrainedTokenizerBase
|
|
|
|
from vllm.benchmarks.datasets import CustomDataset
|
|
from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def hf_tokenizer() -> PreTrainedTokenizerBase:
|
|
# Use a small, commonly available tokenizer
|
|
return AutoTokenizer.from_pretrained("openai-community/gpt2")
|
|
|
|
|
|
text_content = """
|
|
Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
|
|
incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
|
|
exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
|
|
Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat
|
|
nulla pariatur. Excepteur sint occaecat cupidatat non proident,
|
|
sunt in culpa qui officia deserunt mollit anim id est laborum.
|
|
"""
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_create_txt_slices_jsonl(
|
|
hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
|
|
) -> None:
|
|
"""Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset."""
|
|
txt_path = tmp_path / "input.txt"
|
|
jsonl_path = tmp_path / "input.txt.jsonl"
|
|
|
|
txt_path.write_text(text_content)
|
|
|
|
create_txt_slices_jsonl(
|
|
input_path=str(txt_path),
|
|
output_path=str(jsonl_path),
|
|
tokenizer_name="openai-community/gpt2",
|
|
num_prompts=10,
|
|
input_len=10,
|
|
output_len=10,
|
|
)
|
|
|
|
# Verify the JSONL file is valid and has the expected structure
|
|
records = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
|
|
|
|
assert len(records) == 10
|
|
for record in records:
|
|
assert "prompt" in record
|
|
assert "output_tokens" in record
|
|
assert isinstance(record["prompt"], str)
|
|
assert record["output_tokens"] == 10
|
|
|
|
# Verify the JSONL file can be loaded by CustomDataset
|
|
dataset = CustomDataset(dataset_path=str(jsonl_path))
|
|
samples = dataset.sample(
|
|
tokenizer=hf_tokenizer,
|
|
num_requests=10,
|
|
output_len=10,
|
|
skip_chat_template=True,
|
|
)
|
|
|
|
assert len(samples) == 10
|
|
assert all(sample.expected_output_len == 10 for sample in samples)
|