@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 195 lines 4.88 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 from transformers import AutoTokenizer, PreTrainedTokenizerBase
6
7 from tests.reasoning.utils import run_reasoning_extraction
8 from vllm.reasoning import ReasoningParser, ReasoningParserManager
9
10 parser_name = "minimax_m2_append_think"
11 end_token = "</think>"
12
13 # MiniMax M2 model path
14 REASONING_MODEL_NAME = "MiniMaxAI/MiniMax-M2"
15
16
17 @pytest.fixture(scope="module")
18 def minimax_m2_tokenizer() -> PreTrainedTokenizerBase:
19 return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
20
21
22 # =============================================================================
23 # MiniMaxM2AppendThinkReasoningParser behavior:
24 # - Prepends <think> to the beginning of the output
25 # - Does NOT separate reasoning and content
26 # - Returns everything as content (with <think> prepended)
27 # - reasoning is always None
28 #
29 # This parser is used when you want to keep the raw output with <think> added
30 # =============================================================================
31
32 # Case: simple output with end token
33 SIMPLE_OUTPUT = {
34 "output": "This is reasoning</think>This is response",
35 "reasoning": None,
36 "content": "<think>This is reasoning</think>This is response",
37 "is_reasoning_end": True,
38 }
39
40 # Case: output without end token (reasoning in progress)
41 NO_END_TOKEN = {
42 "output": "This is reasoning in progress",
43 "reasoning": None,
44 "content": "<think>This is reasoning in progress",
45 "is_reasoning_end": False,
46 }
47
48 # Case: only end token
49 ONLY_END_TOKEN = {
50 "output": "</think>This is response",
51 "reasoning": None,
52 "content": "<think></think>This is response",
53 "is_reasoning_end": True,
54 }
55
56 # Case: multiple lines
57 MULTIPLE_LINES = {
58 "output": "Line 1\nLine 2</think>Response 1\nResponse 2",
59 "reasoning": None,
60 "content": "<think>Line 1\nLine 2</think>Response 1\nResponse 2",
61 "is_reasoning_end": True,
62 }
63
64 # Case: empty output (non-streaming prepends <think>)
65 EMPTY = {
66 "output": "",
67 "reasoning": None,
68 "content": "<think>",
69 "is_reasoning_end": False,
70 }
71
72 # Case: empty output streaming (no tokens = no output)
73 EMPTY_STREAMING = {
74 "output": "",
75 "reasoning": None,
76 "content": None,
77 "is_reasoning_end": False,
78 }
79
80 # Case: special characters
81 SPECIAL_CHARS = {
82 "output": "Let me think... 1+1=2</think>Yes!",
83 "reasoning": None,
84 "content": "<think>Let me think... 1+1=2</think>Yes!",
85 "is_reasoning_end": True,
86 }
87
88 # Case: code in output
89 CODE_OUTPUT = {
90 "output": "```python\nprint('hi')\n```</think>Here's the code.",
91 "reasoning": None,
92 "content": "<think>```python\nprint('hi')\n```</think>Here's the code.",
93 "is_reasoning_end": True,
94 }
95
96 TEST_CASES = [
97 pytest.param(
98 False,
99 SIMPLE_OUTPUT,
100 id="simple_output",
101 ),
102 pytest.param(
103 True,
104 SIMPLE_OUTPUT,
105 id="simple_output_streaming",
106 ),
107 pytest.param(
108 False,
109 NO_END_TOKEN,
110 id="no_end_token",
111 ),
112 pytest.param(
113 True,
114 NO_END_TOKEN,
115 id="no_end_token_streaming",
116 ),
117 pytest.param(
118 False,
119 ONLY_END_TOKEN,
120 id="only_end_token",
121 ),
122 pytest.param(
123 True,
124 ONLY_END_TOKEN,
125 id="only_end_token_streaming",
126 ),
127 pytest.param(
128 False,
129 MULTIPLE_LINES,
130 id="multiple_lines",
131 ),
132 pytest.param(
133 True,
134 MULTIPLE_LINES,
135 id="multiple_lines_streaming",
136 ),
137 pytest.param(
138 False,
139 EMPTY,
140 id="empty",
141 ),
142 pytest.param(
143 True,
144 EMPTY_STREAMING,
145 id="empty_streaming",
146 ),
147 pytest.param(
148 False,
149 SPECIAL_CHARS,
150 id="special_chars",
151 ),
152 pytest.param(
153 True,
154 SPECIAL_CHARS,
155 id="special_chars_streaming",
156 ),
157 pytest.param(
158 False,
159 CODE_OUTPUT,
160 id="code_output",
161 ),
162 pytest.param(
163 True,
164 CODE_OUTPUT,
165 id="code_output_streaming",
166 ),
167 ]
168
169
170 @pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
171 def test_reasoning(
172 streaming: bool,
173 param_dict: dict,
174 minimax_m2_tokenizer,
175 ):
176 output = minimax_m2_tokenizer.tokenize(param_dict["output"])
177 # decode everything to tokens
178 output_tokens: list[str] = [
179 minimax_m2_tokenizer.convert_tokens_to_string([token]) for token in output
180 ]
181 parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
182 minimax_m2_tokenizer
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"]
191
192 # Test is_reasoning_end
193 output_ids = minimax_m2_tokenizer.convert_tokens_to_ids(output)
194 is_reasoning_end = parser.is_reasoning_end(output_ids)
195 assert is_reasoning_end == param_dict["is_reasoning_end"]