diff options
| -rw-r--r-- | models/llama3/api/test_tokenizer.py | 6 |
1 files changed, 4 insertions, 2 deletions
diff --git a/models/llama3/api/test_tokenizer.py b/models/llama3/api/test_tokenizer.py index c87cc1c..528401a 100644 --- a/models/llama3/api/test_tokenizer.py +++ b/models/llama3/api/test_tokenizer.py @@ -12,7 +12,7 @@ import os from unittest import TestCase from .chat_format import ChatFormat -from .datatypes import SystemMessage, UserMessage +from .datatypes import SystemMessage, ToolPromptFormat, UserMessage from .tokenizer import Tokenizer @@ -49,7 +49,9 @@ class TokenizerTests(TestCase): content="This is a test sentence.", ) self.assertEqual( - self.format.encode_message(message), + self.format.encode_message( + message, tool_prompt_format=ToolPromptFormat.json + )[0], [ 128006, # <|start_header_id|> 882, # "user" |