@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 159 lines 5.92 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Iterable, Sequence
5
6 from vllm.entrypoints.generate.base.protocol import DeltaMessage
7 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
8 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
9 from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
10 from vllm.tokenizers import TokenizerLike
11 from vllm.tool_parsers.utils import partial_tag_overlap
12
13
14 class K2HorizonReasoningParser(DeepSeekR1ReasoningParser):
15 """Reasoning parser for K2 Horizon."""
16
17 _TOOL_CALLS_START = "<ifm|tool_calls>"
18 _EFFORT_TOKENS: dict[str, tuple[str, str]] = {
19 "high": ("<ifm|think>", "</ifm|think>"),
20 "medium": ("<ifm|think_fast>", "</ifm|think_fast>"),
21 "low": ("<ifm|think_faster>", "</ifm|think_faster>"),
22 }
23
24 def __init__(self, tokenizer: TokenizerLike, *args, **kwargs) -> None:
25 chat_template_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
26 effort = chat_template_kwargs.get("reasoning_effort", "high")
27 if not isinstance(effort, str) or effort not in self._EFFORT_TOKENS:
28 supported = ", ".join(self._EFFORT_TOKENS)
29 raise ValueError(
30 f"Unsupported reasoning_effort {effort!r}. "
31 f"Supported values: {supported}."
32 )
33
34 self._start_token, self._end_token = self._EFFORT_TOKENS[effort]
35 super().__init__(tokenizer, *args, **kwargs)
36 self._tool_calls_start_id = self.vocab.get(self._TOOL_CALLS_START)
37 self._stream_buffer = ""
38 self._stream_at_start = True
39 self._stream_finished = False
40
41 @property
42 def start_token(self) -> str:
43 return self._start_token
44
45 @property
46 def end_token(self) -> str:
47 return self._end_token
48
49 def _without_generated_start(self, model_output: str) -> str:
50 return model_output.removeprefix(self.start_token)
51
52 def _split_model_output(self, model_output: str) -> tuple[str, str | None]:
53 output = self._without_generated_start(model_output)
54 if self.end_token in output:
55 reasoning, _, content = output.partition(self.end_token)
56 return reasoning, content or None
57
58 tool_start = output.find(self._TOOL_CALLS_START)
59 if tool_start != -1:
60 return output[:tool_start], output[tool_start:]
61
62 return "", output or None
63
64 def extract_reasoning(
65 self,
66 model_output: str,
67 request: ChatCompletionRequest | ResponsesRequest,
68 ) -> tuple[str, str | None]:
69 del request
70 return self._split_model_output(model_output)
71
72 def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
73 for token_id in reversed(input_ids):
74 if token_id == self.start_token_id:
75 return False
76 if token_id == self.end_token_id or token_id == self._tool_calls_start_id:
77 return True
78 return False
79
80 def is_reasoning_end_streaming(
81 self,
82 input_ids: Sequence[int],
83 delta_ids: Iterable[int],
84 ) -> bool:
85 del input_ids
86 return self._stream_finished or super().is_reasoning_end_streaming(
87 (), delta_ids
88 )
89
90 def extract_content_ids(self, input_ids: list[int]) -> list[int]:
91 if self.end_token_id in input_ids:
92 return input_ids[input_ids.index(self.end_token_id) + 1 :]
93 if (
94 self._tool_calls_start_id is not None
95 and self._tool_calls_start_id in input_ids
96 ):
97 return input_ids[input_ids.index(self._tool_calls_start_id) :]
98 return input_ids if self._stream_finished else []
99
100 def _strip_stream_start(self) -> bool:
101 if not self._stream_at_start:
102 return True
103 if self.start_token.startswith(self._stream_buffer):
104 if self._stream_buffer != self.start_token:
105 return False
106 self._stream_buffer = ""
107 elif self._stream_buffer.startswith(self.start_token):
108 self._stream_buffer = self._stream_buffer[len(self.start_token) :]
109 self._stream_at_start = False
110 return True
111
112 def extract_reasoning_streaming(
113 self,
114 previous_text: str,
115 current_text: str,
116 delta_text: str,
117 previous_token_ids: Sequence[int],
118 current_token_ids: Sequence[int],
119 delta_token_ids: Sequence[int],
120 ) -> DeltaMessage | None:
121 del (
122 previous_text,
123 current_text,
124 previous_token_ids,
125 current_token_ids,
126 delta_token_ids,
127 )
128
129 if self._stream_finished:
130 return DeltaMessage(content=delta_text) if delta_text else None
131
132 self._stream_buffer += delta_text
133 if not self._strip_stream_start():
134 return None
135
136 explicit_end = self._stream_buffer.find(self.end_token)
137 tool_start = self._stream_buffer.find(self._TOOL_CALLS_START)
138 if explicit_end != -1:
139 reasoning = self._stream_buffer[:explicit_end]
140 content = self._stream_buffer[explicit_end + len(self.end_token) :]
141 elif tool_start != -1:
142 reasoning = self._stream_buffer[:tool_start]
143 content = self._stream_buffer[tool_start:]
144 else:
145 overlap = max(
146 partial_tag_overlap(self._stream_buffer, self.end_token),
147 partial_tag_overlap(self._stream_buffer, self._TOOL_CALLS_START),
148 )
149 sendable_end = len(self._stream_buffer) - overlap
150 reasoning = self._stream_buffer[:sendable_end]
151 self._stream_buffer = self._stream_buffer[sendable_end:]
152 return DeltaMessage(reasoning=reasoning) if reasoning else None
153
154 self._stream_buffer = ""
155 self._stream_finished = True
156 return DeltaMessage(
157 reasoning=reasoning,
158 content=content or None,
159 )