@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 190 lines 6.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Make sure bad_words works.
4
5 Run `pytest tests/samplers/test_no_bad_words.py`.
6
7 """
8
9 from transformers import AutoTokenizer
10
11 from tests.conftest import VllmRunner
12 from vllm import SamplingParams
13
14
15 def _generate(
16 llm: VllmRunner,
17 prompt: str,
18 num_prompt_tokens: int,
19 temperature: float = 0,
20 bad_words: list[str] | None = None,
21 ) -> list[int]:
22 sampling_params = SamplingParams(
23 temperature=temperature,
24 bad_words=bad_words,
25 )
26
27 # [([output_token_ids, ], [output_text, ]), ]
28 output = llm.generate([prompt], sampling_params=sampling_params)
29
30 output_token_ids = output[0][0][0][num_prompt_tokens:]
31 # [0] first (and only) request output
32 # [0] token_ids (not text)
33 # [0] first (and only) output completion
34
35 return output_token_ids
36
37
38 class TestOneTokenBadWord:
39 MODEL = "hmellor/tiny-random-LlamaForCausalLM"
40
41 PROMPT = "How old are "
42 TARGET_TOKEN = "mn"
43
44 def setup_method(self, method):
45 self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL)
46
47 self.num_prompt_tokens = len(self._encode(self.PROMPT))
48 self.target_token_id = self._encode(
49 self.TARGET_TOKEN, add_special_tokens=False
50 )[0]
51
52 def test_one_token_bad_word(self, vllm_runner):
53 with vllm_runner(self.MODEL) as llm:
54 output_token_ids = self._generate(llm)
55 assert output_token_ids[0] == self.target_token_id
56
57 output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN])
58 assert self.target_token_id not in output_token_ids
59
60 def _generate(
61 self, llm: VllmRunner, bad_words: list[str] | None = None
62 ) -> list[int]:
63 return _generate(
64 llm=llm,
65 prompt=self.PROMPT,
66 num_prompt_tokens=self.num_prompt_tokens,
67 bad_words=bad_words,
68 )
69
70 def _encode(self, prompt: str, add_special_tokens: bool = True) -> list[int]:
71 return self.tokenizer(prompt, add_special_tokens=add_special_tokens).input_ids
72
73
74 class TestTwoTokenBadWord:
75 # Another model (with a different tokenizer behaviour)
76 MODEL = "distilbert/distilgpt2"
77
78 PROMPT = "How old are you? I am 10"
79 TARGET_TOKEN1 = "years"
80 TARGET_TOKEN2 = "old"
81 NEIGHBOUR_TOKEN2 = "older"
82
83 def setup_method(self, method):
84 self.tokenizer = AutoTokenizer.from_pretrained(
85 self.MODEL, add_prefix_space=True
86 )
87
88 self.num_prompt_tokens = len(self._encode(self.PROMPT))
89 self.target_token_id1 = self._encode(
90 self.TARGET_TOKEN1, add_special_tokens=False
91 )[0]
92 self.target_token_id2 = self._encode(
93 self.TARGET_TOKEN2, add_special_tokens=False
94 )[0]
95 self.neighbour_token_id2 = self._encode(
96 self.NEIGHBOUR_TOKEN2, add_special_tokens=False
97 )[0]
98
99 def test_two_token_bad_word(self, vllm_runner):
100 with vllm_runner(self.MODEL, dtype="half") as llm:
101 output_token_ids = self._generate(llm)
102 assert output_token_ids[:2] == [
103 self.target_token_id1,
104 self.target_token_id2,
105 ]
106
107 output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN1])
108 assert self.target_token_id1 not in output_token_ids
109
110 output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN2])
111 assert output_token_ids[0] == self.target_token_id1
112 assert self.target_token_id2 not in output_token_ids
113
114 output_token_ids = self._generate(
115 llm, bad_words=[f"{self.TARGET_TOKEN1} {self.TARGET_TOKEN2}"]
116 )
117 assert output_token_ids[0] == self.target_token_id1
118 assert output_token_ids[:2] != [
119 self.target_token_id1,
120 self.target_token_id2,
121 ]
122 assert not self._contains(
123 output_token_ids, [self.target_token_id1, self.target_token_id2]
124 )
125 # Model dependent behaviour
126 assert output_token_ids[:2] == [
127 self.target_token_id1,
128 self.neighbour_token_id2,
129 ]
130
131 output_token_ids = self._generate(
132 llm,
133 bad_words=[
134 f"{self.TARGET_TOKEN1} {self.TARGET_TOKEN2}",
135 f"{self.TARGET_TOKEN1} {self.NEIGHBOUR_TOKEN2}",
136 ],
137 )
138 assert output_token_ids[0] == self.target_token_id1
139 assert output_token_ids[:2] != [
140 self.target_token_id1,
141 self.target_token_id2,
142 ]
143 assert not self._contains(
144 output_token_ids, [self.target_token_id1, self.target_token_id2]
145 )
146 assert output_token_ids[:2] != [
147 self.target_token_id1,
148 self.neighbour_token_id2,
149 ]
150 assert not self._contains(
151 output_token_ids, [self.target_token_id1, self.neighbour_token_id2]
152 )
153 assert (self.target_token_id2 in output_token_ids) or (
154 self.neighbour_token_id2 in output_token_ids
155 )
156
157 def _generate(
158 self, llm: VllmRunner, bad_words: list[str] | None = None
159 ) -> list[int]:
160 return _generate(
161 llm=llm,
162 prompt=self.PROMPT,
163 num_prompt_tokens=self.num_prompt_tokens,
164 bad_words=bad_words,
165 )
166
167 @staticmethod
168 def _contains(sequence: list[int], subsequence: list[int]) -> bool:
169 searched = False
170
171 for start in range(len(sequence)):
172 end = start + len(subsequence)
173 current_subsequence = sequence[start:end]
174
175 if len(current_subsequence) < len(subsequence):
176 continue
177
178 searched = True
179
180 assert len(current_subsequence) == len(subsequence)
181
182 if current_subsequence == subsequence:
183 return True
184
185 assert searched, "All subsequences did not match in length..."
186
187 return False
188
189 def _encode(self, prompt: str, add_special_tokens: bool = True) -> list[int]:
190 return self.tokenizer(prompt, add_special_tokens=add_special_tokens).input_ids