main
py 644 lines 19.7 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 (
12 ChatCompletionRequest,
13 )
14 from vllm.parser.abstract_parser import DelegatingParser
15 from vllm.parser.engine.registered_adapters import (
16 GraniteThinkingParserReasoningAdapter,
17 )
18 from vllm.reasoning import ReasoningParser, ReasoningParserManager
19
20 parser_name = "granite_thinking_parser"
21
22
23 class ReasoningCase(TypedDict):
24 output: str
25 reasoning: str | None
26 content: str | None
27
28
29 class GraniteThinkingTokenizer:
30 def __init__(self):
31 self._vocab = {
32 "<think>": 1,
33 "</think>": 2,
34 }
35 self._inv_vocab = {v: k for k, v in self._vocab.items()}
36 self._pattern = re.compile(r"(<think>|</think>)")
37
38 def get_vocab(self) -> dict[str, int]:
39 return self._vocab
40
41 def tokenize(self, text: str) -> list[str]:
42 tokens: list[str] = []
43 for part in self._pattern.split(text):
44 if part:
45 tokens.append(part)
46 return tokens
47
48 def convert_tokens_to_string(self, tokens: list[str]) -> str:
49 return "".join(tokens)
50
51 def decode(self, token_ids: list[int]) -> str:
52 return "".join(self._inv_vocab.get(tid, f"<unk:{tid}>") for tid in token_ids)
53
54
55 @pytest.fixture
56 def tokenizer() -> PreTrainedTokenizerBase:
57 return cast(PreTrainedTokenizerBase, GraniteThinkingTokenizer())
58
59
60 # ── Basic reasoning extraction (non-streaming + streaming) ───────────
61
62
63 @pytest.mark.parametrize(
64 "streaming,param_dict",
65 [
66 pytest.param(
67 False,
68 {
69 "output": "<think>reasoning</think>\nHello",
70 "reasoning": "reasoning",
71 "content": "Hello",
72 },
73 id="leading_newline_stripped",
74 ),
75 pytest.param(
76 True,
77 {
78 "output": "<think>reasoning</think>\nHello",
79 "reasoning": "reasoning",
80 "content": "Hello",
81 },
82 id="leading_newline_stripped_streaming",
83 ),
84 pytest.param(
85 False,
86 {
87 "output": "<think>r</think>c",
88 "reasoning": "r",
89 "content": "c",
90 },
91 id="simple_reasoning",
92 ),
93 pytest.param(
94 True,
95 {
96 "output": "<think>r</think>c",
97 "reasoning": "r",
98 "content": "c",
99 },
100 id="simple_reasoning_streaming",
101 ),
102 pytest.param(
103 False,
104 {
105 "output": "This is a reasoning section</think>This is the rest",
106 "reasoning": "This is a reasoning section",
107 "content": "This is the rest",
108 },
109 id="without_start_token",
110 ),
111 pytest.param(
112 True,
113 {
114 "output": "This is a reasoning section</think>This is the rest",
115 "reasoning": "This is a reasoning section",
116 "content": "This is the rest",
117 },
118 id="without_start_token_streaming",
119 ),
120 pytest.param(
121 False,
122 {
123 "output": "<think>This is a reasoning section</think>This is the rest", # noqa: E501
124 "reasoning": "This is a reasoning section",
125 "content": "This is the rest",
126 },
127 id="with_start_token",
128 ),
129 pytest.param(
130 True,
131 {
132 "output": "<think>This is a reasoning section</think>This is the rest", # noqa: E501
133 "reasoning": "This is a reasoning section",
134 "content": "This is the rest",
135 },
136 id="with_start_token_streaming",
137 ),
138 pytest.param(
139 False,
140 {
141 "output": "<think>reasoning</think>\n\n\nHello",
142 "reasoning": "reasoning",
143 "content": "Hello",
144 },
145 id="multiple_leading_newlines_stripped",
146 ),
147 pytest.param(
148 True,
149 {
150 "output": "<think>reasoning</think>\n\n\nHello",
151 "reasoning": "reasoning",
152 "content": "Hello",
153 },
154 id="multiple_leading_newlines_stripped_streaming",
155 ),
156 pytest.param(
157 False,
158 {
159 "output": "<think>line1\nline2</think>\nresult1\nresult2",
160 "reasoning": "line1\nline2",
161 "content": "result1\nresult2",
162 },
163 id="multiline_reasoning_and_content",
164 ),
165 pytest.param(
166 True,
167 {
168 "output": "<think>line1\nline2</think>\nresult1\nresult2",
169 "reasoning": "line1\nline2",
170 "content": "result1\nresult2",
171 },
172 id="multiline_reasoning_and_content_streaming",
173 ),
174 ],
175 )
176 def test_granite_thinking_reasoning(
177 tokenizer: PreTrainedTokenizerBase,
178 streaming: bool,
179 param_dict: ReasoningCase,
180 ):
181 output = tokenizer.tokenize(param_dict["output"])
182 model_output = [tokenizer.convert_tokens_to_string([token]) for token in output]
183 parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
184 tokenizer
185 )
186
187 reasoning, content = run_reasoning_extraction(
188 parser, model_output, streaming=streaming
189 )
190
191 assert reasoning == param_dict["reasoning"]
192 assert content == param_dict["content"]
193
194
195 # ── No content after end token ───────────────────────────────────────
196
197
198 def test_granite_thinking_no_content_after_end_token(
199 tokenizer: PreTrainedTokenizerBase,
200 ):
201 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
202 parser = parser_cls(tokenizer)
203
204 reasoning, content = run_reasoning_extraction(
205 parser,
206 ["<think>", "reasoning", "</think>"],
207 streaming=False,
208 )
209
210 assert reasoning == "reasoning"
211 assert content is None
212
213
214 # ── Whitespace-only content after end token ──────────────────────────
215
216
217 @pytest.mark.parametrize("streaming", [False, True])
218 def test_granite_thinking_whitespace_only_content(
219 tokenizer: PreTrainedTokenizerBase,
220 streaming: bool,
221 ):
222 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
223 parser = parser_cls(tokenizer)
224
225 reasoning, content = run_reasoning_extraction(
226 parser,
227 ["<think>", "reasoning", "</think>", "\n\n"],
228 streaming=streaming,
229 )
230
231 assert reasoning == "reasoning"
232 assert content is None
233
234
235 # ── Unterminated think block ─────────────────────────────────────────
236
237
238 def test_granite_thinking_unterminated_think_block(
239 tokenizer: PreTrainedTokenizerBase,
240 ):
241 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
242 parser = parser_cls(tokenizer)
243 request = ChatCompletionRequest(
244 model="test-model",
245 messages=[],
246 chat_template_kwargs={"enable_thinking": True},
247 )
248
249 reasoning, content = run_reasoning_extraction(
250 parser,
251 ["<think>", "reasoning only"],
252 request=request,
253 streaming=False,
254 )
255
256 assert reasoning == "reasoning only"
257 assert content is None
258
259
260 # ── enable_thinking=False ────────────────────────────────────────────
261
262
263 def test_granite_thinking_disabled_moves_into_content(
264 tokenizer: PreTrainedTokenizerBase,
265 ):
266 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
267 parser = parser_cls(tokenizer)
268 request = ChatCompletionRequest(
269 model="test-model",
270 messages=[],
271 chat_template_kwargs={"enable_thinking": False},
272 )
273
274 reasoning, content = run_reasoning_extraction(
275 parser,
276 ["This is plain content"],
277 request=request,
278 streaming=False,
279 )
280
281 assert reasoning is None
282 assert content == "This is plain content"
283
284
285 # ── enable_thinking=False + leading newline (§6.2 ordering) ─────────
286
287
288 def test_granite_thinking_disabled_with_leading_newline(
289 tokenizer: PreTrainedTokenizerBase,
290 ):
291 # With enable_thinking=False, model output goes through the swap
292 # path: content is initially None (all text classified as
293 # reasoning), so lstrip doesn't fire pre-swap. After the swap,
294 # any leading \n in the original output is preserved. In practice,
295 # enable_thinking=False output has no template-injected \n, so
296 # this is a correctness check, not a realistic scenario.
297 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
298 parser = parser_cls(tokenizer)
299 request = ChatCompletionRequest(
300 model="test-model",
301 messages=[],
302 chat_template_kwargs={"enable_thinking": False},
303 )
304
305 reasoning, content = run_reasoning_extraction(
306 parser,
307 ["\nThis is plain content"],
308 request=request,
309 streaming=False,
310 )
311
312 assert reasoning is None
313 assert content == "\nThis is plain content"
314
315
316 # ── force_nonempty_content=True ──────────────────────────────────────
317
318
319 def test_granite_thinking_force_nonempty_content_moves_into_content(
320 tokenizer: PreTrainedTokenizerBase,
321 ):
322 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
323 parser = parser_cls(tokenizer)
324 request = ChatCompletionRequest(
325 model="test-model",
326 messages=[],
327 chat_template_kwargs={"force_nonempty_content": True},
328 )
329
330 reasoning, content = run_reasoning_extraction(
331 parser,
332 ["<think>", "This is plain content"],
333 request=request,
334 streaming=False,
335 )
336
337 assert reasoning is None
338 assert content == "This is plain content"
339
340
341 def test_granite_thinking_force_nonempty_no_swap_when_newlines_only(
342 tokenizer: PreTrainedTokenizerBase,
343 ):
344 # When </think> IS present and content is newlines-only, lstrip
345 # removes them but the swap should NOT fire — content was present,
346 # just whitespace. Matches the HF plugin behavior.
347 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
348 parser = parser_cls(tokenizer)
349 request = ChatCompletionRequest(
350 model="test-model",
351 messages=[],
352 chat_template_kwargs={"force_nonempty_content": True},
353 )
354
355 reasoning, content = run_reasoning_extraction(
356 parser,
357 ["<think>", "reasoning", "</think>", "\n\n"],
358 request=request,
359 streaming=False,
360 )
361
362 assert reasoning == "reasoning"
363 assert content is None
364
365
366 def test_granite_thinking_force_nonempty_swaps_when_content_absent(
367 tokenizer: PreTrainedTokenizerBase,
368 ):
369 # When </think> IS present but content is truly absent (zero
370 # characters after </think>, e.g. max_tokens cut), swap fires.
371 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
372 parser = parser_cls(tokenizer)
373 request = ChatCompletionRequest(
374 model="test-model",
375 messages=[],
376 chat_template_kwargs={"force_nonempty_content": True},
377 )
378
379 reasoning, content = run_reasoning_extraction(
380 parser,
381 ["<think>", "reasoning", "</think>"],
382 request=request,
383 streaming=False,
384 )
385
386 assert reasoning is None
387 assert content == "reasoning"
388
389
390 def test_granite_thinking_force_nonempty_keeps_real_content(
391 tokenizer: PreTrainedTokenizerBase,
392 ):
393 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
394 parser = parser_cls(tokenizer)
395 request = ChatCompletionRequest(
396 model="test-model",
397 messages=[],
398 chat_template_kwargs={"force_nonempty_content": True},
399 )
400
401 reasoning, content = run_reasoning_extraction(
402 parser,
403 ["<think>", "reasoning here", "</think>", "real answer"],
404 request=request,
405 streaming=False,
406 )
407
408 assert reasoning == "reasoning here"
409 assert content == "real answer"
410
411
412 # ── Truncated reasoning with thinking on ─────────────────────────────
413
414
415 def test_granite_thinking_keeps_truncated_reasoning(
416 tokenizer: PreTrainedTokenizerBase,
417 ):
418 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
419 parser = parser_cls(tokenizer)
420 request = ChatCompletionRequest(
421 model="test-model",
422 messages=[],
423 chat_template_kwargs={"enable_thinking": True},
424 )
425
426 reasoning, content = run_reasoning_extraction(
427 parser,
428 ["This is truncated reasoning"],
429 request=request,
430 streaming=False,
431 )
432
433 assert reasoning == "This is truncated reasoning"
434 assert content is None
435
436
437 # ── DelegatingParser / parse_delta streaming tests ───────────────────
438
439 _SPECIAL_TOKEN_IDS = {"<think>": 1, "</think>": 2}
440
441
442 def _token_id(token: str) -> int:
443 return _SPECIAL_TOKEN_IDS.get(token, 0)
444
445
446 def _make_reasoning_parser(tokenizer):
447 class _GraniteThinkingDelegating(DelegatingParser):
448 reasoning_parser_cls = GraniteThinkingParserReasoningAdapter
449 tool_parser_cls = None
450
451 return _GraniteThinkingDelegating(tokenizer)
452
453
454 def _run_parse_delta(parser, tokenizer, text, request):
455 tokens = tokenizer.tokenize(text)
456 reasoning_parts: list[str] = []
457 content_parts: list[str] = []
458 for i, token in enumerate(tokens):
459 delta = parser.parse_delta(
460 delta_text=token,
461 delta_token_ids=[_token_id(token)],
462 request=request,
463 prompt_token_ids=[] if i == 0 else None,
464 finished=(i == len(tokens) - 1),
465 )
466 if delta is None:
467 continue
468 if delta.reasoning:
469 reasoning_parts.append(delta.reasoning)
470 if delta.content:
471 content_parts.append(delta.content)
472 return "".join(reasoning_parts), "".join(content_parts)
473
474
475 def test_granite_thinking_streaming_enable_thinking_false(
476 tokenizer: PreTrainedTokenizerBase,
477 ):
478 # With enable_thinking=False, the parser (constructed without
479 # kwargs) starts in REASONING state. All text streams as reasoning
480 # AND is duplicated into content via the streaming fallback —
481 # matching NemotronV3 behavior.
482 request = ChatCompletionRequest(
483 model="test-model",
484 messages=[],
485 chat_template_kwargs={"enable_thinking": False},
486 )
487 parser = _make_reasoning_parser(tokenizer)
488
489 reasoning, content = _run_parse_delta(
490 parser, tokenizer, "This is plain content", request
491 )
492
493 assert reasoning == "This is plain content"
494 assert content == "This is plain content"
495
496
497 def test_granite_thinking_streaming_strips_leading_newline(
498 tokenizer: PreTrainedTokenizerBase,
499 ):
500 request = ChatCompletionRequest(
501 model="test-model",
502 messages=[],
503 )
504 parser = _make_reasoning_parser(tokenizer)
505
506 reasoning, content = _run_parse_delta(
507 parser, tokenizer, "<think>reason</think>\nHello", request
508 )
509
510 assert reasoning == "reason"
511 assert content == "Hello"
512
513
514 def test_granite_thinking_streaming_promotes_reasoning_to_content(
515 tokenizer: PreTrainedTokenizerBase,
516 ):
517 request = ChatCompletionRequest(
518 model="test-model",
519 messages=[],
520 chat_template_kwargs={"force_nonempty_content": True},
521 )
522 parser = _make_reasoning_parser(tokenizer)
523
524 reasoning, content = _run_parse_delta(parser, tokenizer, "<think>4", request)
525
526 assert reasoning == "4"
527 assert content == "4"
528
529
530 def test_granite_thinking_streaming_no_promotion_with_real_content(
531 tokenizer: PreTrainedTokenizerBase,
532 ):
533 request = ChatCompletionRequest(
534 model="test-model",
535 messages=[],
536 chat_template_kwargs={"force_nonempty_content": True},
537 )
538 parser = _make_reasoning_parser(tokenizer)
539
540 reasoning, content = _run_parse_delta(
541 parser,
542 tokenizer,
543 "<think>reason</think>real answer",
544 request,
545 )
546
547 assert reasoning == "reason"
548 assert content == "real answer"
549
550
551 def test_granite_thinking_streaming_no_promotion_without_opt_in(
552 tokenizer: PreTrainedTokenizerBase,
553 ):
554 request = ChatCompletionRequest(model="test-model", messages=[])
555 parser = _make_reasoning_parser(tokenizer)
556
557 reasoning, content = _run_parse_delta(parser, tokenizer, "<think>4", request)
558
559 assert reasoning == "4"
560 assert content == ""
561
562
563 # ── Empty think block ────────────────────────────────────────────────
564
565
566 @pytest.mark.parametrize("streaming", [False, True])
567 def test_granite_thinking_empty_think_block(
568 tokenizer: PreTrainedTokenizerBase,
569 streaming: bool,
570 ):
571 # <think></think>\nHello — empty reasoning, content after newline.
572 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
573 parser = parser_cls(tokenizer)
574
575 reasoning, content = run_reasoning_extraction(
576 parser,
577 ["<think>", "</think>", "\n", "Hello"],
578 streaming=streaming,
579 )
580
581 assert reasoning is None or reasoning == ""
582 assert content == "Hello"
583
584
585 # ── Chunking-independence streaming test ─────────────────────────────
586
587
588 @pytest.mark.parametrize(
589 "text,expected_reasoning,expected_content",
590 [
591 ("<think>r</think>\nHello", "r", "Hello"),
592 ("<think></think>\nHello", "", "Hello"),
593 ],
594 )
595 def test_granite_thinking_streaming_chunking_independent(
596 tokenizer: PreTrainedTokenizerBase,
597 text: str,
598 expected_reasoning: str,
599 expected_content: str,
600 ):
601 # Verify the same input produces identical results regardless
602 # of how it's chunked into streaming deltas. Uses the
603 # DelegatingParser path (_run_parse_delta) which matches the
604 # real serving flow.
605 tokens = tokenizer.tokenize(text)
606 request = ChatCompletionRequest(model="test-model", messages=[])
607
608 chunk_patterns = [
609 tokens,
610 # Split at the </think> boundary
611 [
612 tokenizer.convert_tokens_to_string(tokens[: tokens.index("</think>") + 1]),
613 tokenizer.convert_tokens_to_string(tokens[tokens.index("</think>") + 1 :]),
614 ],
615 ]
616
617 for chunks in chunk_patterns:
618 parser = _make_reasoning_parser(tokenizer)
619 reasoning_parts: list[str] = []
620 content_parts: list[str] = []
621 all_tokens = []
622 for chunk in chunks:
623 chunk_tokens = tokenizer.tokenize(chunk)
624 all_tokens.extend(chunk_tokens)
625 for i, token in enumerate(all_tokens):
626 delta = parser.parse_delta(
627 delta_text=token,
628 delta_token_ids=[_token_id(token)],
629 request=request,
630 prompt_token_ids=[] if i == 0 else None,
631 finished=(i == len(all_tokens) - 1),
632 )
633 if delta is None:
634 continue
635 if delta.reasoning:
636 reasoning_parts.append(delta.reasoning)
637 if delta.content:
638 content_parts.append(delta.content)
639 reasoning = "".join(reasoning_parts)
640 content = "".join(content_parts)
641 assert reasoning == expected_reasoning or (
642 not expected_reasoning and not reasoning
643 ), f"reasoning mismatch with chunks={chunks}"
644 assert content == expected_content, f"content mismatch with chunks={chunks}"