@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 113 lines 3.77 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Nemotron V3 parser.
4
5 The Nemotron 3 Super model uses the same tool call and reasoning
6 format as Qwen3 (``<think>``/``</think>`` + ``<tool_call>`` XML).
7 This config reuses :func:`qwen3_config` with a distinct name.
8
9 When ``enable_thinking=False`` or ``force_nonempty_content=True`` and
10 content is empty, reasoning and content are swapped.
11 """
12
13 from __future__ import annotations
14
15 import dataclasses
16 import functools
17 from typing import TYPE_CHECKING
18
19 from vllm.parser.qwen3 import CHATML_TURN_BOUNDARIES, Qwen3Parser, qwen3_config
20
21 if TYPE_CHECKING:
22 from vllm.entrypoints.generate.base.protocol import DeltaMessage
23 from vllm.entrypoints.openai.chat_completion.protocol import (
24 ChatCompletionRequest,
25 )
26 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
27 from vllm.parser.engine.parser_engine import SemanticEvent
28 from vllm.parser.engine.parser_engine_config import ParserEngineConfig
29 from vllm.tokenizers import TokenizerLike
30 from vllm.tool_parsers.abstract_tool_parser import Tool
31
32
33 @functools.cache
34 def nemotron_v3_config(thinking: bool = True) -> ParserEngineConfig:
35 return dataclasses.replace(
36 qwen3_config(thinking=thinking, turn_boundary_tokens=CHATML_TURN_BOUNDARIES),
37 name="nemotron_v3",
38 strip_trailing_reasoning_whitespace=True,
39 )
40
41
42 class NemotronV3Parser(Qwen3Parser):
43 """Nemotron V3 parser: same format as Qwen3, with Nemotron-specific
44 behavior: when ``enable_thinking=False`` or
45 ``force_nonempty_content=True`` and content is empty, swaps
46 reasoning and content.
47 """
48
49 def __init__(
50 self,
51 tokenizer: TokenizerLike,
52 tools: list[Tool] | None = None,
53 **kwargs,
54 ) -> None:
55 chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
56 thinking = chat_kwargs.get("enable_thinking", True)
57 super().__init__(
58 tokenizer,
59 tools,
60 parser_engine_config=nemotron_v3_config(thinking=thinking),
61 **kwargs,
62 )
63 self._streamed_reasoning: list[str] = []
64
65 def _reset(self, initial_state=None) -> None:
66 super()._reset(initial_state=initial_state)
67 self._streamed_reasoning = []
68
69 def _events_to_delta(
70 self,
71 events: list[SemanticEvent],
72 finished: bool = False,
73 ) -> DeltaMessage | None:
74 delta = super()._events_to_delta(events, finished=finished)
75 if delta is not None and delta.reasoning is not None:
76 self._streamed_reasoning.append(delta.reasoning)
77 return delta
78
79 @staticmethod
80 def _should_force_content(
81 request: ChatCompletionRequest | ResponsesRequest,
82 ) -> bool:
83 chat_template_kwargs = getattr(request, "chat_template_kwargs", None)
84 return bool(
85 chat_template_kwargs
86 and (
87 chat_template_kwargs.get("enable_thinking") is False
88 or chat_template_kwargs.get("force_nonempty_content") is True
89 )
90 )
91
92 def get_streaming_fallback_content(
93 self,
94 text: str,
95 request: ChatCompletionRequest | ResponsesRequest,
96 ) -> str | None:
97 if not self._should_force_content(request):
98 return None
99 return "".join(self._streamed_reasoning) or None
100
101 def extract_reasoning(
102 self,
103 model_output: str,
104 request: ChatCompletionRequest | ResponsesRequest,
105 ) -> tuple[str | None, str | None]:
106 reasoning, content = super().extract_reasoning(model_output, request)
107
108 if self._should_force_content(request) and (
109 content is None or not content.strip()
110 ):
111 reasoning, content = content, reasoning
112
113 return reasoning, content