Fix test_tokenizer.py

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"