| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | from transformers import AutoTokenizer |
| 6 | |
| 7 | from vllm import SamplingParams |
| 8 | from vllm.v1.engine import EngineCoreRequest |
| 9 | from vllm.v1.engine.detokenizer import FastIncrementalDetokenizer |
| 10 | |
| 11 | PROMPT = "Hello, my name is Lee, and I'm a student in the " + "college of engineering" |
| 12 | |
| 13 | |
| 14 | @pytest.mark.parametrize( |
| 15 | "min_tokens,stop,truth", |
| 16 | [ |
| 17 | (0, None, " is Lee, and I'm a student in the college of engineering"), |
| 18 | (0, "e", " is L"), |
| 19 | (5, "e", " is Lee, and I'm a stud"), |
| 20 | ], |
| 21 | ) |
| 22 | def test_min_tokens_with_stop(min_tokens: int, stop: str, truth: str): |
| 23 | """Test for a specific min_tokens and stop. |
| 24 | |
| 25 | See https://github.com/vllm-project/vllm/pull/22014 |
| 26 | """ |
| 27 | tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m") |
| 28 | all_prompt_ids = tokenizer(PROMPT, add_special_tokens=False).input_ids |
| 29 | |
| 30 | # The prompt is "Hello, my name is" |
| 31 | prompt_token_ids = all_prompt_ids[:4] |
| 32 | params = SamplingParams( |
| 33 | stop=stop, |
| 34 | min_tokens=min_tokens, |
| 35 | ) |
| 36 | request = EngineCoreRequest( |
| 37 | request_id="", |
| 38 | prompt_token_ids=prompt_token_ids, |
| 39 | mm_features=None, |
| 40 | sampling_params=params, |
| 41 | pooling_params=None, |
| 42 | arrival_time=0.0, |
| 43 | lora_request=None, |
| 44 | cache_salt=None, |
| 45 | data_parallel_rank=None, |
| 46 | ) |
| 47 | |
| 48 | detokenizer = FastIncrementalDetokenizer(tokenizer, request) |
| 49 | |
| 50 | detokenizer.update(all_prompt_ids[4:], False) |
| 51 | assert detokenizer.output_text == truth |