@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 168 lines 4.59 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
6
7 from tests.reasoning.utils import run_reasoning_extraction
8 from vllm.reasoning import ReasoningParser, ReasoningParserManager
9
10 parser_name = "hunyuan_a13b"
11 START_REASONING = "<think>\n"
12 START_RESPONSE = "\n</think>\n<answer>\n"
13 END_RESPONSE = "\n</answer>"
14
15 NO_REASONING_QUICK_THOUGHT = {
16 "output": f"{START_REASONING}{START_RESPONSE}This is the rest{END_RESPONSE}", # noqa: E501
17 "reasoning": None,
18 "content": "This is the rest",
19 }
20
21 SIMPLE_REASONING = {
22 "output": f"{START_REASONING}This is a reasoning section{START_RESPONSE}This is the rest{END_RESPONSE}", # noqa: E501
23 "reasoning": "This is a reasoning section",
24 "content": "This is the rest",
25 }
26 COMPLETE_REASONING = {
27 "output": f"{START_REASONING}This is a reasoning section{START_RESPONSE}",
28 "reasoning": "This is a reasoning section",
29 "content": None,
30 }
31
32 COMPLETE_REASONING_WITH_SYMBOL = {
33 "output": f"{START_REASONING}This is a reasoning section!{START_RESPONSE}",
34 "reasoning": "This is a reasoning section!",
35 "content": None,
36 }
37 NO_REASONING = {
38 "output": "This is content",
39 "reasoning": None,
40 "content": "This is content",
41 }
42 MULTIPLE_LINES = {
43 "output": f"{START_REASONING}This\nThat{START_RESPONSE}This is the rest\nThat",
44 "reasoning": "This\nThat",
45 "content": "This is the rest\nThat",
46 }
47 REASONING_WITH_THINK = {
48 "output": f"{START_REASONING}This is a reasoning section{START_RESPONSE}This is the rest", # noqa: E501
49 "reasoning": "This is a reasoning section",
50 "content": "This is the rest",
51 }
52 COMPLETE_REASONING_WITH_THINK = {
53 "output": f"{START_REASONING}This is a reasoning section{START_RESPONSE}",
54 "reasoning": "This is a reasoning section",
55 "content": None,
56 }
57 MULTIPLE_LINES_WITH_THINK = {
58 "output": f"{START_REASONING}This\nThat{START_RESPONSE}This is the rest\nThat",
59 "reasoning": "This\nThat",
60 "content": "This is the rest\nThat",
61 }
62
63 TEST_CASES = [
64 pytest.param(
65 False,
66 SIMPLE_REASONING,
67 id="simple_reasoning",
68 ),
69 pytest.param(
70 False,
71 COMPLETE_REASONING,
72 id="complete_reasoning",
73 ),
74 pytest.param(
75 False,
76 COMPLETE_REASONING_WITH_SYMBOL,
77 id="complete_reasoning_with_symbol",
78 ),
79 pytest.param(
80 False,
81 NO_REASONING,
82 id="no_reasoning",
83 ),
84 pytest.param(False, NO_REASONING_QUICK_THOUGHT, id="no_reasoning_quick"),
85 pytest.param(
86 False,
87 MULTIPLE_LINES,
88 id="multiple_lines",
89 ),
90 pytest.param(
91 False,
92 REASONING_WITH_THINK,
93 id="reasoning_with_think",
94 ),
95 pytest.param(
96 False,
97 COMPLETE_REASONING_WITH_THINK,
98 id="complete_reasoning_with_think",
99 ),
100 pytest.param(
101 False,
102 MULTIPLE_LINES_WITH_THINK,
103 id="multiple_lines_with_think",
104 ),
105 pytest.param(
106 True,
107 SIMPLE_REASONING,
108 id="simple_reasoning_streaming",
109 ),
110 pytest.param(
111 True,
112 COMPLETE_REASONING,
113 id="complete_reasoning_streaming",
114 ),
115 pytest.param(
116 True,
117 NO_REASONING,
118 id="no_reasoning_streaming",
119 ),
120 pytest.param(True, NO_REASONING_QUICK_THOUGHT, id="no_reasoning_quick_stream"),
121 pytest.param(
122 True,
123 MULTIPLE_LINES,
124 id="multiple_lines_streaming",
125 ),
126 pytest.param(
127 True,
128 REASONING_WITH_THINK,
129 id="reasoning_with_think_streaming",
130 ),
131 pytest.param(
132 True,
133 COMPLETE_REASONING_WITH_THINK,
134 id="complete_reasoning_with_think_streaming",
135 ),
136 pytest.param(
137 True,
138 MULTIPLE_LINES_WITH_THINK,
139 id="multiple_lines_with_think_streaming",
140 ),
141 ]
142
143 # Global tokenizer initialization to avoid repeated loading
144 tokenizer = AutoTokenizer.from_pretrained(
145 "tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
146 )
147
148
149 @pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
150 def test_reasoning(
151 streaming: bool,
152 param_dict: dict,
153 ):
154 output = tokenizer.tokenize(param_dict["output"])
155 # decode everything to tokens
156 output_tokens: list[str] = [
157 tokenizer.convert_tokens_to_string([token]) for token in output
158 ]
159 parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
160 tokenizer
161 )
162
163 reasoning, content = run_reasoning_extraction(
164 parser, output_tokens, streaming=streaming
165 )
166
167 assert reasoning == param_dict["reasoning"]
168 assert content == param_dict["content"]