@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 293 lines 8.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import TypedDict, cast
5
6 import pytest
7 import regex as re
8 from transformers import PreTrainedTokenizerBase
9
10 from tests.reasoning.utils import run_reasoning_extraction
11 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
12 from vllm.parser.abstract_parser import DelegatingParser
13 from vllm.parser.engine.registered_adapters import NemotronV3ParserReasoningAdapter
14 from vllm.reasoning import ReasoningParser, ReasoningParserManager
15
16 parser_name = "nemotron_v3"
17
18
19 class ReasoningCase(TypedDict):
20 output: str
21 reasoning: str | None
22 content: str | None
23
24
25 class FakeNemotronTokenizer:
26 def __init__(self):
27 self._vocab = {
28 "<think>": 1,
29 "</think>": 2,
30 }
31 self._inv_vocab = {v: k for k, v in self._vocab.items()}
32 self._pattern = re.compile(r"(<think>|</think>)")
33
34 def get_vocab(self) -> dict[str, int]:
35 return self._vocab
36
37 def tokenize(self, text: str) -> list[str]:
38 tokens: list[str] = []
39 for part in self._pattern.split(text):
40 if part:
41 tokens.append(part)
42 return tokens
43
44 def convert_tokens_to_string(self, tokens: list[str]) -> str:
45 return "".join(tokens)
46
47 def decode(self, token_ids: list[int]) -> str:
48 return "".join(self._inv_vocab.get(tid, f"<unk:{tid}>") for tid in token_ids)
49
50
51 @pytest.fixture
52 def tokenizer() -> PreTrainedTokenizerBase:
53 return cast(PreTrainedTokenizerBase, FakeNemotronTokenizer())
54
55
56 @pytest.mark.parametrize(
57 "streaming,param_dict",
58 [
59 pytest.param(
60 False,
61 {
62 "output": "This is a reasoning section</think>This is the rest",
63 "reasoning": "This is a reasoning section",
64 "content": "This is the rest",
65 },
66 id="without_start_token",
67 ),
68 pytest.param(
69 True,
70 {
71 "output": "This is a reasoning section</think>This is the rest",
72 "reasoning": "This is a reasoning section",
73 "content": "This is the rest",
74 },
75 id="without_start_token_streaming",
76 ),
77 pytest.param(
78 False,
79 {
80 "output": "<think>This is a reasoning section</think>This is the rest",
81 "reasoning": "This is a reasoning section",
82 "content": "This is the rest",
83 },
84 id="with_start_token",
85 ),
86 pytest.param(
87 True,
88 {
89 "output": "<think>This is a reasoning section</think>This is the rest",
90 "reasoning": "This is a reasoning section",
91 "content": "This is the rest",
92 },
93 id="with_start_token_streaming",
94 ),
95 ],
96 )
97 def test_nemotron_v3_reasoning(
98 tokenizer: PreTrainedTokenizerBase,
99 streaming: bool,
100 param_dict: ReasoningCase,
101 ):
102 output = tokenizer.tokenize(param_dict["output"])
103 model_output = [tokenizer.convert_tokens_to_string([token]) for token in output]
104 parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
105 tokenizer
106 )
107
108 reasoning, content = run_reasoning_extraction(
109 parser, model_output, streaming=streaming
110 )
111
112 assert reasoning == param_dict["reasoning"]
113 assert content == param_dict["content"]
114
115
116 def test_nemotron_v3_without_thinking_moves_into_content(
117 tokenizer: PreTrainedTokenizerBase,
118 ):
119 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
120 parser = parser_cls(tokenizer)
121 request = ChatCompletionRequest(
122 model="test-model",
123 messages=[],
124 chat_template_kwargs={"enable_thinking": False},
125 )
126
127 reasoning, content = run_reasoning_extraction(
128 parser,
129 ["This is plain content"],
130 request=request,
131 streaming=False,
132 )
133
134 # No real content followed the reasoning, so the trace is moved into
135 # content (reasoning left empty) — matching main's behavior.
136 assert reasoning is None
137 assert content == "This is plain content"
138
139
140 def test_nemotron_v3_force_nonempty_content_moves_into_content(
141 tokenizer: PreTrainedTokenizerBase,
142 ):
143 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
144 parser = parser_cls(tokenizer)
145 request = ChatCompletionRequest(
146 model="test-model",
147 messages=[],
148 chat_template_kwargs={"force_nonempty_content": True},
149 )
150
151 reasoning, content = run_reasoning_extraction(
152 parser,
153 ["<think>This is plain content"],
154 request=request,
155 streaming=False,
156 )
157
158 assert reasoning is None
159 assert content == "This is plain content"
160
161
162 def test_nemotron_v3_force_nonempty_keeps_real_content(
163 tokenizer: PreTrainedTokenizerBase,
164 ):
165 # When real content follows the closing tag nothing is promoted: the
166 # content after </think> is returned as-is and reasoning stays separate.
167 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
168 parser = parser_cls(tokenizer)
169 request = ChatCompletionRequest(
170 model="test-model",
171 messages=[],
172 chat_template_kwargs={"force_nonempty_content": True},
173 )
174
175 reasoning, content = run_reasoning_extraction(
176 parser,
177 ["<think>reasoning here</think>real answer"],
178 request=request,
179 streaming=False,
180 )
181
182 assert reasoning == "reasoning here"
183 assert content == "real answer"
184
185
186 def test_nemotron_v3_with_thinking_keeps_truncated_reasoning(
187 tokenizer: PreTrainedTokenizerBase,
188 ):
189 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
190 parser = parser_cls(tokenizer)
191 request = ChatCompletionRequest(
192 model="test-model",
193 messages=[],
194 chat_template_kwargs={"enable_thinking": True},
195 )
196
197 reasoning, content = run_reasoning_extraction(
198 parser,
199 ["This is truncated reasoning"],
200 request=request,
201 streaming=False,
202 )
203
204 assert reasoning == "This is truncated reasoning"
205 assert content is None
206
207
208 _SPECIAL_TOKEN_IDS = {"<think>": 1, "</think>": 2}
209
210
211 def _token_id(token: str) -> int:
212 # Only the think markers need stable ids; everything else is non-special.
213 return _SPECIAL_TOKEN_IDS.get(token, 0)
214
215
216 def _make_reasoning_parser(tokenizer):
217 class _NemotronParser(DelegatingParser):
218 reasoning_parser_cls = NemotronV3ParserReasoningAdapter
219 tool_parser_cls = None
220
221 return _NemotronParser(tokenizer)
222
223
224 def _run_parse_delta(parser, tokenizer, text, request):
225 tokens = tokenizer.tokenize(text)
226 reasoning_parts: list[str] = []
227 content_parts: list[str] = []
228 for i, token in enumerate(tokens):
229 delta = parser.parse_delta(
230 delta_text=token,
231 delta_token_ids=[_token_id(token)],
232 request=request,
233 prompt_token_ids=[] if i == 0 else None,
234 finished=(i == len(tokens) - 1),
235 )
236 if delta is None:
237 continue
238 if delta.reasoning:
239 reasoning_parts.append(delta.reasoning)
240 if delta.content:
241 content_parts.append(delta.content)
242 return "".join(reasoning_parts), "".join(content_parts)
243
244
245 def test_nemotron_v3_streaming_promotes_reasoning_to_content(
246 tokenizer: PreTrainedTokenizerBase,
247 ):
248 # Model never closes <think>: reasoning streams normally AND is duplicated
249 # into content on the terminal delta.
250 request = ChatCompletionRequest(
251 model="test-model",
252 messages=[],
253 chat_template_kwargs={"force_nonempty_content": True},
254 )
255 parser = _make_reasoning_parser(tokenizer)
256
257 reasoning, content = _run_parse_delta(parser, tokenizer, "<think>4", request)
258
259 assert reasoning == "4"
260 assert content == "4"
261
262
263 def test_nemotron_v3_streaming_no_promotion_with_real_content(
264 tokenizer: PreTrainedTokenizerBase,
265 ):
266 request = ChatCompletionRequest(
267 model="test-model",
268 messages=[],
269 chat_template_kwargs={"force_nonempty_content": True},
270 )
271 parser = _make_reasoning_parser(tokenizer)
272
273 reasoning, content = _run_parse_delta(
274 parser, tokenizer, "<think>reason</think>real answer", request
275 )
276
277 # Real content followed </think>, so nothing is duplicated.
278 assert reasoning == "reason"
279 assert content == "real answer"
280
281
282 def test_nemotron_v3_streaming_no_promotion_without_opt_in(
283 tokenizer: PreTrainedTokenizerBase,
284 ):
285 # Without enable_thinking=False / force_nonempty_content the fallback must
286 # stay disabled: the response stays reasoning-only, content empty.
287 request = ChatCompletionRequest(model="test-model", messages=[])
288 parser = _make_reasoning_parser(tokenizer)
289
290 reasoning, content = _run_parse_delta(parser, tokenizer, "<think>4", request)
291
292 assert reasoning == "4"
293 assert content == ""