@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 4.83 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 = "olmo3"
11 START_REASONING = "<think>"
12 END_REASONING = "</think>"
13
14 NO_REASONING = {
15 "output": f"{START_REASONING}{END_REASONING}No thoughts, head empty!",
16 "reasoning": None,
17 "content": "No thoughts, head empty!",
18 }
19
20 NO_REASONING_WITH_NEWLINE = {
21 "output": f"{START_REASONING}\n{END_REASONING}\n\nNo thoughts, head empty!",
22 "reasoning": "\n",
23 "content": "\n\nNo thoughts, head empty!",
24 }
25
26 SIMPLE_REASONING = {
27 "output": f"{START_REASONING}This is a reasoning section{END_REASONING}This is the rest", # noqa: E501
28 "reasoning": "This is a reasoning section",
29 "content": "This is the rest",
30 }
31
32 SIMPLE_REASONING_WITH_NEWLINE = {
33 "output": f"{START_REASONING} Look!\n\nI'm thinking...{END_REASONING}\nThis is the rest", # noqa: E501
34 "reasoning": " Look!\n\nI'm thinking...",
35 "content": "\nThis is the rest",
36 }
37
38 SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES = {
39 "output": f"{START_REASONING}\nLook!\nI'm thinking...\n\n{END_REASONING}\n\n\nThis is the rest", # noqa: E501
40 "reasoning": "\nLook!\nI'm thinking...\n\n",
41 "content": "\n\n\nThis is the rest",
42 }
43
44 SIMPLE_REASONING_WITH_TRAILING_SPACE = {
45 "output": f"{START_REASONING}\nLook!\nI'm thinking... {END_REASONING}\nThis is the rest", # noqa: E501
46 "reasoning": "\nLook!\nI'm thinking... ",
47 "content": "\nThis is the rest",
48 }
49
50 NO_REASONING_ONLY_END_THINK = {
51 "output": f"{END_REASONING}\n\nNo thoughts, head empty!",
52 "reasoning": None,
53 "content": "\n\nNo thoughts, head empty!",
54 }
55
56 REASONING_ONLY_END_THINK = {
57 "output": f"The user is asking me not to think.{END_REASONING}No thoughts!",
58 "reasoning": "The user is asking me not to think.",
59 "content": "No thoughts!",
60 }
61
62 TEST_CASES = [
63 pytest.param(
64 False, # not streaming
65 NO_REASONING,
66 id="no_reasoning",
67 ),
68 pytest.param(
69 False, # not streaming
70 NO_REASONING_WITH_NEWLINE,
71 id="no_reasoning_with_newline",
72 ),
73 pytest.param(
74 False, # not streaming
75 SIMPLE_REASONING,
76 id="simple_reasoning",
77 ),
78 pytest.param(
79 False, # not streaming
80 SIMPLE_REASONING_WITH_NEWLINE,
81 id="simple_reasoning_with_newline",
82 ),
83 pytest.param(
84 True, # enable streaming
85 SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES,
86 id="simple_reasoning_with_multiple_newlines",
87 ),
88 pytest.param(
89 False, # not streaming
90 NO_REASONING_ONLY_END_THINK,
91 id="no_reasoning_only_end_think",
92 ),
93 pytest.param(
94 False, # not streaming
95 REASONING_ONLY_END_THINK,
96 id="yes_reasoning_only_end_think",
97 ),
98 pytest.param(
99 True, # enable streaming
100 NO_REASONING,
101 id="no_reasoning_streaming",
102 ),
103 pytest.param(
104 True, # enable streaming
105 NO_REASONING_WITH_NEWLINE,
106 id="no_reasoning_with_newline_streaming",
107 ),
108 pytest.param(
109 True, # enable streaming
110 SIMPLE_REASONING,
111 id="simple_reasoning_streaming",
112 ),
113 pytest.param(
114 True, # enable streaming
115 SIMPLE_REASONING_WITH_NEWLINE,
116 id="simple_reasoning_with_newline_streaming",
117 ),
118 pytest.param(
119 True, # enable streaming
120 SIMPLE_REASONING_WITH_MULTIPLE_NEWLINES,
121 id="simple_reasoning_with_multiple_newlines_streaming",
122 ),
123 pytest.param(
124 True, # enable streaming
125 SIMPLE_REASONING_WITH_TRAILING_SPACE,
126 id="simple_reasoning_with_trailing_space_streaming",
127 ),
128 pytest.param(
129 True, # enable streaming
130 NO_REASONING_ONLY_END_THINK,
131 id="no_reasoning_only_end_think_streaming",
132 ),
133 pytest.param(
134 True, # enable streaming
135 REASONING_ONLY_END_THINK,
136 id="yes_reasoning_only_end_think_streaming",
137 ),
138 ]
139
140 # Global tokenizer initialization to avoid repeated loading
141 tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-7B-Think")
142
143
144 @pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
145 def test_reasoning(
146 streaming: bool,
147 param_dict: dict[str, str],
148 ):
149 output = tokenizer.tokenize(param_dict["output"])
150
151 # decode everything to tokens
152 model_output: list[str] = [
153 tokenizer.convert_tokens_to_string([token]) for token in output
154 ]
155 parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
156 parser: ReasoningParser = parser_cls(tokenizer)
157
158 reasoning, content = run_reasoning_extraction(
159 reasoning_parser=parser, model_output=model_output, streaming=streaming
160 )
161
162 assert reasoning == param_dict["reasoning"]
163 assert content == param_dict["content"]