| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import json |
| 5 | from pathlib import Path |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.benchmarks.datasets import get_samples |
| 10 | |
| 11 | |
| 12 | class _RecordingTokenizer: |
| 13 | """Minimal tokenizer stub that records the kwargs forwarded to |
| 14 | apply_chat_template, so we can assert chat_template_kwargs propagation |
| 15 | without loading a real model/template.""" |
| 16 | |
| 17 | def __init__(self) -> None: |
| 18 | self.captured_kwargs: dict | None = None |
| 19 | self.chat_template = "dummy-template" |
| 20 | |
| 21 | def apply_chat_template( |
| 22 | self, |
| 23 | conversation, |
| 24 | add_generation_prompt: bool = True, |
| 25 | tokenize: bool = False, |
| 26 | **kwargs, |
| 27 | ) -> str: |
| 28 | self.captured_kwargs = kwargs |
| 29 | return conversation[0]["content"] |
| 30 | |
| 31 | def __call__(self, text: str): |
| 32 | return argparse.Namespace(input_ids=list(range(len(text.split())))) |
| 33 | |
| 34 | |
| 35 | def _args(dataset_path: str, chat_template_kwargs) -> argparse.Namespace: |
| 36 | return argparse.Namespace( |
| 37 | dataset_name="custom", |
| 38 | dataset_path=dataset_path, |
| 39 | disable_shuffle=True, |
| 40 | num_prompts=1, |
| 41 | custom_output_len=32, |
| 42 | skip_chat_template=False, |
| 43 | chat_template_kwargs=chat_template_kwargs, |
| 44 | no_oversample=False, |
| 45 | seed=0, |
| 46 | request_id_prefix="", |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | def _write_one(path: Path) -> None: |
| 51 | path.write_text(json.dumps({"prompt": "hello world"}) + "\n") |
| 52 | |
| 53 | |
| 54 | @pytest.mark.benchmark |
| 55 | def test_chat_template_kwargs_forwarded(tmp_path: Path) -> None: |
| 56 | """--chat-template-kwargs must reach the client-side apply_chat_template.""" |
| 57 | jsonl = tmp_path / "data.jsonl" |
| 58 | _write_one(jsonl) |
| 59 | |
| 60 | tok = _RecordingTokenizer() |
| 61 | get_samples(_args(str(jsonl), {"thinking": True}), tok) |
| 62 | |
| 63 | assert tok.captured_kwargs == {"thinking": True} |
| 64 | |
| 65 | |
| 66 | @pytest.mark.benchmark |
| 67 | def test_chat_template_kwargs_default_is_noop(tmp_path: Path) -> None: |
| 68 | """When not provided, no extra kwargs are passed (existing behavior).""" |
| 69 | jsonl = tmp_path / "data.jsonl" |
| 70 | _write_one(jsonl) |
| 71 | |
| 72 | tok = _RecordingTokenizer() |
| 73 | get_samples(_args(str(jsonl), None), tok) |
| 74 | |
| 75 | assert tok.captured_kwargs == {} |