| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from vllm import LLM, SamplingParams |
| 7 | from vllm.exceptions import VLLMValidationError |
| 8 | |
| 9 | MODEL = "hmellor/tiny-random-LlamaForCausalLM" |
| 10 | PROMPT = "Hello my name is Robert and I" |
| 11 | |
| 12 | |
| 13 | @pytest.fixture(scope="module") |
| 14 | def llm() -> LLM: |
| 15 | return LLM(MODEL, enforce_eager=True) |
| 16 | |
| 17 | |
| 18 | def test_n_gt_1(llm): |
| 19 | """ParallelSampling is supported.""" |
| 20 | params = SamplingParams(n=3) |
| 21 | outputs = llm.generate(PROMPT, params) |
| 22 | assert len(outputs[0].outputs) == 3 |
| 23 | |
| 24 | |
| 25 | def test_penalties(llm): |
| 26 | """Check that we do not get errors if applied.""" |
| 27 | params = SamplingParams( |
| 28 | temperature=1.2, |
| 29 | presence_penalty=1.2, |
| 30 | frequency_penalty=1.2, |
| 31 | repetition_penalty=1.2, |
| 32 | min_p=0.5, |
| 33 | top_p=0.5, |
| 34 | top_k=3, |
| 35 | ) |
| 36 | _ = llm.generate(PROMPT, params) |
| 37 | |
| 38 | |
| 39 | def test_stop(llm): |
| 40 | """Check that we respect the stop words.""" |
| 41 | output = llm.generate(PROMPT, SamplingParams(temperature=0)) |
| 42 | split_text = output[0].outputs[0].text.split() |
| 43 | |
| 44 | STOP_IDX = 5 |
| 45 | params = SamplingParams(temperature=0, stop=split_text[STOP_IDX]) |
| 46 | output = llm.generate(PROMPT, params) |
| 47 | new_split_text = output[0].outputs[0].text.split() |
| 48 | |
| 49 | # Output should not contain the stop word. |
| 50 | assert len(new_split_text) == STOP_IDX |
| 51 | |
| 52 | params = SamplingParams( |
| 53 | temperature=0, stop=split_text[STOP_IDX], include_stop_str_in_output=True |
| 54 | ) |
| 55 | output = llm.generate(PROMPT, params) |
| 56 | new_split_text = output[0].outputs[0].text.split() |
| 57 | |
| 58 | # Output should contain the stop word. |
| 59 | assert len(new_split_text) == STOP_IDX + 1 |
| 60 | |
| 61 | |
| 62 | def test_stop_token_ids(llm): |
| 63 | """Check that we respect the stop token ids.""" |
| 64 | output = llm.generate(PROMPT, SamplingParams(temperature=0)) |
| 65 | |
| 66 | stop_token_id_0 = output[0].outputs[0].token_ids[5] |
| 67 | stop_token_id_1 = output[0].outputs[0].token_ids[6] |
| 68 | |
| 69 | stop_token_ids = [stop_token_id_1, stop_token_id_0] |
| 70 | params = SamplingParams(temperature=0, stop_token_ids=stop_token_ids) |
| 71 | output = llm.generate(PROMPT, params) |
| 72 | assert output[0].outputs[0].token_ids[-1] == stop_token_id_0 |
| 73 | |
| 74 | stop_token_ids = [stop_token_id_0, stop_token_id_1] |
| 75 | params = SamplingParams(temperature=0, stop_token_ids=stop_token_ids) |
| 76 | output = llm.generate(PROMPT, params) |
| 77 | assert output[0].outputs[0].token_ids[-1] == stop_token_id_0 |
| 78 | |
| 79 | |
| 80 | def test_detokenize_false(llm): |
| 81 | """Check that detokenize=False option works.""" |
| 82 | output = llm.generate(PROMPT, SamplingParams(detokenize=False)) |
| 83 | assert len(output[0].outputs[0].token_ids) > 0 |
| 84 | assert len(output[0].outputs[0].text) == 0 |
| 85 | |
| 86 | output = llm.generate( |
| 87 | PROMPT, SamplingParams(detokenize=False, logprobs=3, prompt_logprobs=3) |
| 88 | ) |
| 89 | assert len(output[0].outputs[0].token_ids) > 0 |
| 90 | assert len(output[0].outputs[0].text) == 0 |
| 91 | |
| 92 | prompt_logprobs = output[0].prompt_logprobs |
| 93 | sampled_logprobs = output[0].outputs[0].logprobs |
| 94 | assert len(prompt_logprobs) > 1 |
| 95 | assert len(sampled_logprobs) > 1 |
| 96 | for all_logprobs in (prompt_logprobs[1:], sampled_logprobs): |
| 97 | for logprobs in all_logprobs: |
| 98 | assert 3 <= len(logprobs) <= 4 |
| 99 | assert all(lp.decoded_token is None for lp in logprobs.values()) |
| 100 | |
| 101 | |
| 102 | def test_bad_words(llm): |
| 103 | """Check that we respect bad words.""" |
| 104 | tokenizer = llm.get_tokenizer() |
| 105 | |
| 106 | def contains_bad_word(text: str, tokens: list[int], bad_word: str) -> bool: |
| 107 | """Check if word appears in BOTH text and token sequence.""" |
| 108 | if bad_word not in text: |
| 109 | return False |
| 110 | |
| 111 | for add_prefix_space in [False, True]: |
| 112 | prefix = " " if add_prefix_space else "" |
| 113 | bad_words_token = tokenizer.encode( |
| 114 | prefix + bad_word.lstrip(), add_special_tokens=False |
| 115 | ) |
| 116 | if not bad_words_token: |
| 117 | continue |
| 118 | for i in range(len(tokens) - len(bad_words_token) + 1): |
| 119 | if tokens[i : i + len(bad_words_token)] == bad_words_token: |
| 120 | return True |
| 121 | return False |
| 122 | |
| 123 | output = llm.generate(PROMPT, SamplingParams(temperature=0)) |
| 124 | split_text = output[0].outputs[0].text.split() |
| 125 | |
| 126 | bad_words_1 = " ".join(split_text[:2]) |
| 127 | params = SamplingParams(temperature=0, bad_words=[bad_words_1]) |
| 128 | output = llm.generate(PROMPT, params) |
| 129 | new_text = output[0].outputs[0].text |
| 130 | new_tokens = output[0].outputs[0].token_ids |
| 131 | assert not contains_bad_word(new_text, new_tokens, bad_words_1) |
| 132 | |
| 133 | bad_words_2 = new_text.split()[-1] |
| 134 | params = SamplingParams(temperature=0, bad_words=[bad_words_1, bad_words_2]) |
| 135 | output = llm.generate(PROMPT, params) |
| 136 | new_text = output[0].outputs[0].text |
| 137 | new_tokens = output[0].outputs[0].token_ids |
| 138 | assert not contains_bad_word(new_text, new_tokens, bad_words_1) |
| 139 | assert not contains_bad_word(new_text, new_tokens, bad_words_2) |
| 140 | |
| 141 | |
| 142 | def test_allowed_token_ids(llm): |
| 143 | """Check that we can use allowed_token_ids.""" |
| 144 | TOKEN_ID = 10 |
| 145 | allowed_token_ids = [TOKEN_ID] |
| 146 | output = llm.generate(PROMPT, SamplingParams(allowed_token_ids=allowed_token_ids)) |
| 147 | assert output[0].outputs[0].token_ids[-1] == TOKEN_ID |
| 148 | |
| 149 | # Each single-token allowlist must force that token (kernel used to drop some). |
| 150 | for token_id in (1, 5, 100, 500, 2518, 9834, 31999): |
| 151 | output = llm.generate( |
| 152 | PROMPT, |
| 153 | SamplingParams(temperature=0, max_tokens=1, allowed_token_ids=[token_id]), |
| 154 | ) |
| 155 | assert output[0].outputs[0].token_ids[-1] == token_id |
| 156 | |
| 157 | # Reject empty allowed_token_ids. |
| 158 | with pytest.raises(VLLMValidationError): |
| 159 | _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[])) |
| 160 | |
| 161 | # Reject negative token id. |
| 162 | with pytest.raises(VLLMValidationError): |
| 163 | _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[-1])) |
| 164 | |
| 165 | # Reject out of vocabulary. |
| 166 | with pytest.raises(VLLMValidationError): |
| 167 | _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[10000000])) |
| 168 | |
| 169 | |
| 170 | def test_seed(llm): |
| 171 | """Check that seed impacts randomness.""" |
| 172 | out_1 = llm.generate(PROMPT, SamplingParams(seed=42)) |
| 173 | out_2 = llm.generate(PROMPT, SamplingParams(seed=42)) |
| 174 | out_3 = llm.generate(PROMPT, SamplingParams(seed=43)) |
| 175 | |
| 176 | assert out_1[0].outputs[0].text == out_2[0].outputs[0].text |
| 177 | assert out_1[0].outputs[0].text != out_3[0].outputs[0].text |