@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 190 lines 5.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 from transformers import AutoTokenizer, PreTrainedTokenizerBase
5
6 from tests.reasoning.utils import run_reasoning_extraction
7 from vllm.reasoning import ReasoningParser, ReasoningParserManager
8 from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
9 from vllm.reasoning.deepseek_v3_reasoning_parser import (
10 DeepSeekV3ReasoningWithThinkingParser as Holo2ReasoningParser,
11 )
12 from vllm.reasoning.identity_reasoning_parser import IdentityReasoningParser
13
14 REASONING_MODEL_NAME = "HCompany/Holo2-4B"
15
16
17 @pytest.fixture(scope="module")
18 def tokenizer() -> PreTrainedTokenizerBase:
19 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
20
21
22 @pytest.mark.parametrize(
23 "thinking,expected_parser_type",
24 [
25 (True, DeepSeekR1ReasoningParser),
26 (False, IdentityReasoningParser),
27 ],
28 )
29 def test_parser_selection(tokenizer, thinking, expected_parser_type):
30 parser = Holo2ReasoningParser(
31 tokenizer,
32 chat_template_kwargs={
33 "thinking": thinking,
34 },
35 )
36
37 assert isinstance(parser._parser, expected_parser_type)
38
39
40 def test_holo2_default_parser_is_deepseekr1(tokenizer):
41 parser = Holo2ReasoningParser(tokenizer)
42
43 assert isinstance(parser._parser, DeepSeekR1ReasoningParser)
44
45
46 def test_holo2_supports_structured_output(tokenizer):
47 # Structured output manager uses the reasoning parser to check if the
48 # reasoning content is ended before applying the grammar. The main function
49 # used is is_reasoning_end. This test checks if the parser is able to
50 # correctly identify the end of the reasoning content.
51
52 # important to not pass chat_template_kwargs here as it is done in the
53 # StructuredOutputManager
54 parser = Holo2ReasoningParser(tokenizer)
55
56 end_token_id = tokenizer.encode("</think>", add_special_tokens=False)[0]
57
58 assert parser.is_reasoning_end([1, 2, 4, end_token_id])
59 assert not parser.is_reasoning_end([1, 2, 4])
60 assert parser.is_reasoning_end([1, 2, 4, end_token_id, 5])
61
62
63 # thinking is True, non-streaming
64 WITH_THINK = {
65 "output": "This is a reasoning section</think>This is the rest",
66 "reasoning": "This is a reasoning section",
67 "content": "This is the rest",
68 }
69 # thinking is True, streaming
70 WITH_THINK_STREAM = {
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 # thinking is False, non-streaming
76 THINKING_DISABLED = {
77 "output": "This is the rest",
78 "reasoning": None,
79 "content": "This is the rest",
80 }
81 # thinking is False, streaming
82 THINKING_DISABLED_STREAM = {
83 "output": "This is the rest",
84 "reasoning": None,
85 "content": "This is the rest",
86 }
87 # thinking is False but the model output </think>, non-streaming
88 THINKING_DISABLED_WITH_CLOSE_TAG = {
89 "output": "</think>This is the rest",
90 "reasoning": None,
91 "content": "</think>This is the rest",
92 }
93 # thinking is False but the model output </think>, streaming
94 THINKING_DISABLED_WITH_CLOSE_TAG_STREAM = {
95 "output": "some text</think>This is the rest",
96 "reasoning": None,
97 "content": "some text</think>This is the rest",
98 }
99 COMPLETE_REASONING = {
100 "output": "This is a reasoning section</think>",
101 "reasoning": "This is a reasoning section",
102 "content": None,
103 }
104
105 TEST_CASES = [
106 pytest.param(
107 False,
108 WITH_THINK,
109 None,
110 id="with_think",
111 ),
112 pytest.param(
113 True,
114 WITH_THINK_STREAM,
115 None,
116 id="with_think_stream",
117 ),
118 pytest.param(
119 False,
120 WITH_THINK,
121 {"thinking": True},
122 id="with_think_enabled",
123 ),
124 pytest.param(
125 True,
126 WITH_THINK_STREAM,
127 {"thinking": True},
128 id="with_think_stream_enabled",
129 ),
130 pytest.param(
131 False,
132 THINKING_DISABLED,
133 {"thinking": False},
134 id="thinking_disabled",
135 ),
136 pytest.param(
137 True,
138 THINKING_DISABLED_STREAM,
139 {"thinking": False},
140 id="thinking_disabled_stream",
141 ),
142 pytest.param(
143 False,
144 THINKING_DISABLED_WITH_CLOSE_TAG,
145 {"thinking": False},
146 id="thinking_disabled_with_close_tag",
147 ),
148 pytest.param(
149 True,
150 THINKING_DISABLED_WITH_CLOSE_TAG_STREAM,
151 {"thinking": False},
152 id="thinking_disabled_with_close_tag_stream",
153 ),
154 pytest.param(
155 False,
156 COMPLETE_REASONING,
157 None,
158 id="complete_reasoning",
159 ),
160 pytest.param(
161 True,
162 COMPLETE_REASONING,
163 None,
164 id="complete_reasoning_stream",
165 ),
166 ]
167
168
169 @pytest.mark.parametrize("streaming, param_dict, chat_template_kwargs", TEST_CASES)
170 def test_reasoning(
171 streaming: bool,
172 param_dict: dict,
173 chat_template_kwargs: dict | None,
174 tokenizer,
175 ):
176 output = tokenizer.tokenize(param_dict["output"])
177 output_tokens: list[str] = [
178 tokenizer.convert_tokens_to_string([token]) for token in output
179 ]
180 parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser("holo2")(
181 tokenizer,
182 chat_template_kwargs=chat_template_kwargs,
183 )
184
185 reasoning, content = run_reasoning_extraction(
186 parser, output_tokens, streaming=streaming
187 )
188
189 assert reasoning == param_dict["reasoning"]
190 assert content == param_dict["content"]