@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 204 lines 5.53 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"
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 # MiniMax M2 specific behavior:
24 # - Model does NOT generate <think> start token
25 # - Model only generates </think> end token
26 # - All content before </think> is reasoning
27 # - All content after </think> is the actual response (content)
28 # =============================================================================
29
30 # Case: reasoning + end token + content (typical case)
31 SIMPLE_REASONING = {
32 "output": "This is a reasoning section</think>This is the rest",
33 "reasoning": "This is a reasoning section",
34 "content": "This is the rest",
35 "is_reasoning_end": True,
36 }
37
38 # Case: reasoning + end token only (no content after)
39 COMPLETE_REASONING = {
40 "output": "This is a reasoning section</think>",
41 "reasoning": "This is a reasoning section",
42 "content": None,
43 "is_reasoning_end": True,
44 }
45
46 # Case: no end token yet (streaming in progress, all is reasoning)
47 NO_END_TOKEN = {
48 "output": "This is reasoning in progress",
49 "reasoning": "This is reasoning in progress",
50 "content": None,
51 "is_reasoning_end": False,
52 }
53
54 # Case: multiple lines of reasoning
55 MULTIPLE_LINES = {
56 "output": "First line\nSecond line</think>Response first line\nResponse second",
57 "reasoning": "First line\nSecond line",
58 "content": "Response first line\nResponse second",
59 "is_reasoning_end": True,
60 }
61
62 # Case: only end token streaming (reasoning is None because it's just the token)
63 SHORTEST_REASONING_STREAMING = {
64 "output": "</think>This is the response",
65 "reasoning": None,
66 "content": "This is the response",
67 "is_reasoning_end": True,
68 }
69
70 # Case: empty streaming
71 EMPTY_STREAMING = {
72 "output": "",
73 "reasoning": None,
74 "content": None,
75 "is_reasoning_end": False,
76 }
77
78 # Case: long reasoning with special characters
79 SPECIAL_CHARS = {
80 "output": "Let me think... 1+1=2, right?</think>Yes, 1+1=2.",
81 "reasoning": "Let me think... 1+1=2, right?",
82 "content": "Yes, 1+1=2.",
83 "is_reasoning_end": True,
84 }
85
86 # Case: reasoning with code blocks
87 CODE_IN_REASONING = {
88 "output": "```python\nprint('hello')\n```</think>Here is the code.",
89 "reasoning": "```python\nprint('hello')\n```",
90 "content": "Here is the code.",
91 "is_reasoning_end": True,
92 }
93
94 TEST_CASES = [
95 # Core cases: no start token (MiniMax M2 actual behavior)
96 pytest.param(
97 False,
98 SIMPLE_REASONING,
99 id="simple_reasoning",
100 ),
101 pytest.param(
102 True,
103 SIMPLE_REASONING,
104 id="simple_reasoning_streaming",
105 ),
106 pytest.param(
107 False,
108 COMPLETE_REASONING,
109 id="complete_reasoning",
110 ),
111 pytest.param(
112 True,
113 COMPLETE_REASONING,
114 id="complete_reasoning_streaming",
115 ),
116 pytest.param(
117 False,
118 NO_END_TOKEN,
119 id="no_end_token",
120 ),
121 pytest.param(
122 True,
123 NO_END_TOKEN,
124 id="no_end_token_streaming",
125 ),
126 pytest.param(
127 False,
128 MULTIPLE_LINES,
129 id="multiple_lines",
130 ),
131 pytest.param(
132 True,
133 MULTIPLE_LINES,
134 id="multiple_lines_streaming",
135 ),
136 pytest.param(
137 True,
138 SHORTEST_REASONING_STREAMING,
139 id="shortest_reasoning_streaming",
140 ),
141 pytest.param(
142 True,
143 EMPTY_STREAMING,
144 id="empty_streaming",
145 ),
146 pytest.param(
147 False,
148 SPECIAL_CHARS,
149 id="special_chars",
150 ),
151 pytest.param(
152 True,
153 SPECIAL_CHARS,
154 id="special_chars_streaming",
155 ),
156 pytest.param(
157 False,
158 CODE_IN_REASONING,
159 id="code_in_reasoning",
160 ),
161 pytest.param(
162 True,
163 CODE_IN_REASONING,
164 id="code_in_reasoning_streaming",
165 ),
166 ]
167
168
169 @pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
170 def test_reasoning(
171 streaming: bool,
172 param_dict: dict,
173 minimax_m2_tokenizer,
174 ):
175 output = minimax_m2_tokenizer.tokenize(param_dict["output"])
176 # decode everything to tokens
177 output_tokens: list[str] = [
178 minimax_m2_tokenizer.convert_tokens_to_string([token]) for token in output
179 ]
180 parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
181 minimax_m2_tokenizer
182 )
183
184 reasoning, content = run_reasoning_extraction(
185 parser, output_tokens, streaming=streaming
186 )
187
188 assert reasoning == param_dict["reasoning"]
189 assert content == param_dict["content"]
190
191 # Test is_reasoning_end
192 output_ids = minimax_m2_tokenizer.convert_tokens_to_ids(output)
193 is_reasoning_end = parser.is_reasoning_end(output_ids)
194 assert is_reasoning_end == param_dict["is_reasoning_end"]
195
196 # Test extract_content
197 if param_dict["content"] is not None:
198 content_ids = parser.extract_content_ids(output_ids)
199 assert content_ids == minimax_m2_tokenizer.convert_tokens_to_ids(
200 minimax_m2_tokenizer.tokenize(param_dict["content"])
201 )
202 else:
203 content_ids = parser.extract_content_ids(output)
204 assert content_ids == []