Skip to content

Commit 932a093

Browse files
committed
fix mypy problem
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
1 parent a90055f commit 932a093

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

vllm/entrypoints/openai/tool_parsers/llama_tool_parser.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import partial_json_parser
77
from partial_json_parser.core.options import Allow
8+
from transformers import PreTrainedTokenizerBase
89

910
from vllm.entrypoints.openai.protocol import (DeltaFunctionCall, DeltaMessage,
1011
DeltaToolCall,
@@ -14,7 +15,6 @@
1415
ToolParser)
1516
from vllm.entrypoints.openai.tool_parsers.utils import find_common_prefix
1617
from vllm.logger import init_logger
17-
from vllm.transformers_utils.tokenizer import AnyTokenizer
1818
from vllm.utils import random_uuid
1919

2020
logger = init_logger(__name__)
@@ -49,7 +49,7 @@ class Llama3JsonToolParser(ToolParser):
4949
Used when --enable-auto-tool-choice --tool-call-parser mistral are all set
5050
"""
5151

52-
def __init__(self, tokenizer: AnyTokenizer):
52+
def __init__(self, tokenizer: PreTrainedTokenizerBase):
5353
super().__init__(tokenizer)
5454

5555
# initialize properties used for state when parsing tool calls in
@@ -60,8 +60,8 @@ def __init__(self, tokenizer: AnyTokenizer):
6060
self.streamed_args_for_tool: List[str] = [
6161
] # map what has been streamed for each tool so far to a list
6262
self.bot_token = "<|python_tag|>"
63-
self.bot_token_id = self.model_tokenizer.encode(
64-
self.bot_token, add_special_tokens=False)[0]
63+
self.bot_token_id = tokenizer.encode(self.bot_token,
64+
add_special_tokens=False)[0]
6565
self.tool_call_regex = re.compile(r"\[{.*?}\]", re.DOTALL)
6666

6767
def extract_tool_calls(self,

0 commit comments

Comments
 (0)