| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test the different finish_reason="stop" situations during generation: |
| 4 | 1. One of the provided stop strings |
| 5 | 2. One of the provided stop tokens |
| 6 | 3. The EOS token |
| 7 | |
| 8 | Run `pytest tests/engine/test_stop_reason.py`. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | import transformers |
| 13 | |
| 14 | from vllm import SamplingParams |
| 15 | |
| 16 | MODEL = "distilbert/distilgpt2" |
| 17 | STOP_STR = "." |
| 18 | SEED = 42 |
| 19 | MAX_TOKENS = 1024 |
| 20 | |
| 21 | |
| 22 | @pytest.fixture |
| 23 | def vllm_model(vllm_runner): |
| 24 | with vllm_runner(MODEL) as vllm_model: |
| 25 | yield vllm_model |
| 26 | |
| 27 | |
| 28 | def test_stop_reason(vllm_model, example_prompts): |
| 29 | tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL) |
| 30 | stop_token_id = tokenizer.convert_tokens_to_ids(STOP_STR) |
| 31 | llm = vllm_model.llm |
| 32 | |
| 33 | # test stop token |
| 34 | outputs = llm.generate( |
| 35 | example_prompts, |
| 36 | sampling_params=SamplingParams( |
| 37 | ignore_eos=True, |
| 38 | seed=SEED, |
| 39 | max_tokens=MAX_TOKENS, |
| 40 | stop_token_ids=[stop_token_id], |
| 41 | ), |
| 42 | ) |
| 43 | for output in outputs: |
| 44 | output = output.outputs[0] |
| 45 | assert output.finish_reason == "stop" |
| 46 | assert output.stop_reason == stop_token_id |
| 47 | |
| 48 | # test stop string |
| 49 | outputs = llm.generate( |
| 50 | example_prompts, |
| 51 | sampling_params=SamplingParams( |
| 52 | ignore_eos=True, seed=SEED, max_tokens=MAX_TOKENS, stop="." |
| 53 | ), |
| 54 | ) |
| 55 | for output in outputs: |
| 56 | output = output.outputs[0] |
| 57 | assert output.finish_reason == "stop" |
| 58 | assert output.stop_reason == STOP_STR |
| 59 | |
| 60 | # test EOS token |
| 61 | outputs = llm.generate( |
| 62 | example_prompts, |
| 63 | sampling_params=SamplingParams(seed=SEED, max_tokens=MAX_TOKENS), |
| 64 | ) |
| 65 | for output in outputs: |
| 66 | output = output.outputs[0] |
| 67 | assert output.finish_reason == "length" or ( |
| 68 | output.finish_reason == "stop" and output.stop_reason is None |
| 69 | ) |