| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | from pathlib import Path |
| 5 | |
| 6 | import pytest |
| 7 | from transformers import AutoTokenizer, PreTrainedTokenizerBase |
| 8 | |
| 9 | from vllm.benchmarks.datasets import CustomDataset |
| 10 | from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl |
| 11 | |
| 12 | |
| 13 | @pytest.fixture(scope="session") |
| 14 | def hf_tokenizer() -> PreTrainedTokenizerBase: |
| 15 | # Use a small, commonly available tokenizer |
| 16 | return AutoTokenizer.from_pretrained("openai-community/gpt2") |
| 17 | |
| 18 | |
| 19 | text_content = """ |
| 20 | Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor |
| 21 | incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud |
| 22 | exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. |
| 23 | Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat |
| 24 | nulla pariatur. Excepteur sint occaecat cupidatat non proident, |
| 25 | sunt in culpa qui officia deserunt mollit anim id est laborum. |
| 26 | """ |
| 27 | |
| 28 | |
| 29 | @pytest.mark.benchmark |
| 30 | def test_create_txt_slices_jsonl( |
| 31 | hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path |
| 32 | ) -> None: |
| 33 | """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset.""" |
| 34 | txt_path = tmp_path / "input.txt" |
| 35 | jsonl_path = tmp_path / "input.txt.jsonl" |
| 36 | |
| 37 | txt_path.write_text(text_content) |
| 38 | |
| 39 | create_txt_slices_jsonl( |
| 40 | input_path=str(txt_path), |
| 41 | output_path=str(jsonl_path), |
| 42 | tokenizer_name="openai-community/gpt2", |
| 43 | num_prompts=10, |
| 44 | input_len=10, |
| 45 | output_len=10, |
| 46 | ) |
| 47 | |
| 48 | # Verify the JSONL file is valid and has the expected structure |
| 49 | records = [json.loads(line) for line in jsonl_path.read_text().splitlines()] |
| 50 | |
| 51 | assert len(records) == 10 |
| 52 | for record in records: |
| 53 | assert "prompt" in record |
| 54 | assert "output_tokens" in record |
| 55 | assert isinstance(record["prompt"], str) |
| 56 | assert record["output_tokens"] == 10 |
| 57 | |
| 58 | # Verify the JSONL file can be loaded by CustomDataset |
| 59 | dataset = CustomDataset(dataset_path=str(jsonl_path)) |
| 60 | samples = dataset.sample( |
| 61 | tokenizer=hf_tokenizer, |
| 62 | num_requests=10, |
| 63 | output_len=10, |
| 64 | skip_chat_template=True, |
| 65 | ) |
| 66 | |
| 67 | assert len(samples) == 10 |
| 68 | assert all(sample.expected_output_len == 10 for sample in samples) |