Skip to content

GitLab

Unverified Commit beb16b2c authored by Travis Johnson's avatar Travis Johnson Committed by GitHub
Browse files

[Bugfix] Handle <|tool_call|> token in granite tool parser (#11039) 


Signed-off-by: default avatarTravis Johnson <tsjohnso@us.ibm.com>
parent fe2e10c7
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
vllm/entrypoints/openai/tool_parsers/granite_tool_parser.py
View file @ beb16b2c
...@@ -35,11 +35,13 @@ class GraniteToolParser(ToolParser): ...@@ -35,11 +35,13 @@ class GraniteToolParser(ToolParser):
def __init__(self, tokenizer: AnyTokenizer): def __init__(self, tokenizer: AnyTokenizer):
super().__init__(tokenizer) super().__init__(tokenizer)
self.bot_token = "<|tool_call|>"
def extract_tool_calls( def extract_tool_calls(
self, model_output: str, self, model_output: str,
request: ChatCompletionRequest) -> ExtractedToolCallInformation: request: ChatCompletionRequest) -> ExtractedToolCallInformation:
stripped = model_output.strip() # remove whitespace and the BOT token if it exists
stripped = model_output.strip().removeprefix(self.bot_token).lstrip()
if not stripped or stripped[0] != '[': if not stripped or stripped[0] != '[':
return ExtractedToolCallInformation(tools_called=False, return ExtractedToolCallInformation(tools_called=False,
tool_calls=[], tool_calls=[],
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment