Upload 2 files

llama_nemotron_nano_generic_tool_calling.jinja

@@ -0,0 +1,51 @@
+{%- if tools %}
+    {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n' -}}
+    {%- if messages[0].role == 'system' -%}
+        {{- messages[0].content + '\n\n' -}}
+    {%- endif -%}
+    {{- '<AVAILABLE_TOOLS>[' -}}
+    {%- for tool in tools -%}
+        {{- (tool.function if tool.function is defined else tool) | tojson -}}{{- ', ' if not loop.last else '' -}}
+    {%- endfor -%}
+    {{- ']</AVAILABLE_TOOLS>' -}}{{- '<|eot_id|>' -}}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n' + messages[0].content + '<|eot_id|>' -}}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages -%}
+    {%- if (message.role == 'user') -%}
+        {{- '<|start_header_id|>user<|end_header_id|>\n\n' + message.content + '<|eot_id|>' -}}
+    {%- elif message.role == 'assistant' -%}
+        {%- set content = message.content -%}
+        {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' + content -}}
+        {%- if message.tool_calls -%}
+            {{- '<TOOLCALL>[' -}}
+            {%- for tool_call in message.tool_calls -%}
+                {%- if tool_call.function -%}
+                    {%- set tool_call = tool_call.function -%}
+                {%- endif -%}
+                {{- '{"name": "' }}{{- tool_call.name }}{{- '", "arguments": ' -}}
+                {%- if tool_call.arguments is string -%}
+                    {{- tool_call.arguments -}}
+                {%- else -%}
+                    {{- tool_call.arguments | tojson -}}
+                {%- endif -%}
+                {{- ', ' if not loop.last else '' -}}
+            {%- endfor -%}
+            {{- ']</TOOLCALL>' -}}
+        {%- endif %}
+        {{- '<|eot_id|>' -}}
+    {%- elif message.role == 'tool' -%}
+        {%- if loop.first or (messages[loop.index0 - 1].role != 'tool') -%}
+            {{- '<|start_header_id|>user<|end_header_id|>\n\n' }}{{- '<TOOL_RESPONSE>[' -}}
+        {%- endif -%}
+        {{- message.content -}}{{- ', ' if not loop.last and (messages[loop.index0 + 1].role == 'tool') else '' -}}
+        {%- if loop.last or (messages[loop.index0 + 1].role != 'tool') -%}
+            {{- ']</TOOL_RESPONSE>' -}}{{- '<|eot_id|>' -}}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' -}}
+{%- endif %}

llama_nemotron_nano_toolcall_parser.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import re
+from collections.abc import Sequence
+from typing import Union
+
+import partial_json_parser
+from partial_json_parser.core.options import Allow
+
+from vllm.entrypoints.openai.protocol import (
+    ChatCompletionRequest,
+    DeltaFunctionCall, DeltaMessage,
+    DeltaToolCall,
+    ExtractedToolCallInformation,
+    FunctionCall,
+    ToolCall,
+)
+from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
+    ToolParser,
+    ToolParserManager,
+)
+from vllm.logger import init_logger
+from vllm.transformers_utils.tokenizer import AnyTokenizer
+from vllm.utils import random_uuid
+
+logger = init_logger(__name__)
+
+
+@ToolParserManager.register_module("llama_nemotron_json")
+class LlamaNemotronJSONToolParser(ToolParser):
+
+    def __init__(self, tokenizer: AnyTokenizer):
+        super().__init__(tokenizer)
+
+        self.current_tool_name_sent: bool = False
+        self.prev_tool_call_arr: list[dict] = []
+        self.current_tool_id: int = -1
+        self.streamed_args_for_tool: list[str] = []
+
+        self.tool_call_start_token: str = "<TOOLCALL>"
+        self.tool_call_end_token: str = "</TOOLCALL>"
+
+        self.tool_call_regex = re.compile(r"<TOOLCALL>(.*?)</TOOLCALL>", re.DOTALL)
+
+    def extract_tool_calls(
+        self,
+        model_output: str,
+        request: ChatCompletionRequest,
+    ) -> ExtractedToolCallInformation:
+
+        if self.tool_call_start_token not in model_output:
+            return ExtractedToolCallInformation(
+                tools_called=False,
+                tool_calls=[],
+                content=model_output,
+            )
+
+        else:
+
+            try:
+                str_tool_calls = self.tool_call_regex.findall(model_output)[0].strip()
+                if not str_tool_calls.startswith("["):
+                    str_tool_calls = "[" + str_tool_calls
+                if not str_tool_calls.endswith("]"):
+                    str_tool_calls = "]" + str_tool_calls
+                json_tool_calls = json.loads(str_tool_calls)
+                tool_calls = []
+                for tool_call in json_tool_calls:
+                    try:
+                        tool_calls.append(ToolCall(
+                            type="function",
+                            function=FunctionCall(
+                                name=tool_call["name"],
+                                arguments=json.dumps(tool_call["arguments"], ensure_ascii=False) \
+                                    if isinstance(tool_call["arguments"], dict) else tool_call["arguments"],
+                            ),
+                        ))
+                    except:
+                        continue
+
+                content = model_output[:model_output.rfind(self.tool_call_start_token)]
+
+                return ExtractedToolCallInformation(
+                    tools_called=True,
+                    tool_calls=tool_calls,
+                    content=content if content else None,
+                )
+
+            except Exception:
+                logger.exception(f"Error in extracting tool call from response. Response: {model_output}")
+                return ExtractedToolCallInformation(
+                    tools_called=False,
+                    tool_calls=[],
+                    content=model_output,
+                )
+
+    def extract_tool_calls_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+        request: ChatCompletionRequest,
+    ) -> Union[DeltaMessage, None]:
+
+        raise NotImplementedError("Tool calling is not supported in streaming mode!")