fix: fixed initial_prompt_tokens calculation

epam · Nov 8, 2024 · c68499f · c68499f
1 parent b727ff1
commit c68499f
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 13 deletions.
diff --git a/aidial_adapter_openai/gpt.py b/aidial_adapter_openai/gpt.py
@@ -23,14 +23,17 @@
 
 
 def plain_text_truncate_prompt(
-    messages: List[dict], max_prompt_tokens: int, tokenizer: PlainTextTokenizer
+    request: dict,
+    messages: List[dict],
+    max_prompt_tokens: int,
+    tokenizer: PlainTextTokenizer,
 ) -> Tuple[List[dict], DiscardedMessages, TruncatedTokens]:
     return truncate_prompt(
         messages=messages,
         message_tokens=tokenizer.tokenize_request_message,
         is_system_message=lambda message: message["role"] == "system",
         max_prompt_tokens=max_prompt_tokens,
-        initial_prompt_tokens=tokenizer.TOKENS_PER_REQUEST,
+        initial_prompt_tokens=tokenizer.tokenize_request(request, []),
     )
 
 
@@ -58,6 +61,7 @@ async def gpt_chat_completion(
 
         request["messages"], discarded_messages, estimated_prompt_tokens = (
             plain_text_truncate_prompt(
+                request=request,
                 messages=cast(List[dict], request["messages"]),
                 max_prompt_tokens=max_prompt_tokens,
                 tokenizer=tokenizer,

diff --git a/aidial_adapter_openai/gpt4_multi_modal/chat_completion.py b/aidial_adapter_openai/gpt4_multi_modal/chat_completion.py
@@ -119,9 +119,9 @@ async def predict_non_stream(
 
 
 def multi_modal_truncate_prompt(
+    request: dict,
     messages: List[MultiModalMessage],
     max_prompt_tokens: int,
-    initial_prompt_tokens: int,
     tokenizer: MultiModalTokenizer,
 ) -> Tuple[List[MultiModalMessage], DiscardedMessages, TruncatedTokens]:
     return truncate_prompt(
@@ -130,7 +130,7 @@ def multi_modal_truncate_prompt(
         is_system_message=lambda message: message.raw_message["role"]
         == "system",
         max_prompt_tokens=max_prompt_tokens,
-        initial_prompt_tokens=initial_prompt_tokens,
+        initial_prompt_tokens=tokenizer.tokenize_request(request, []),
     )
 
 
@@ -218,9 +218,9 @@ async def chat_completion(
     if max_prompt_tokens is not None:
         multi_modal_messages, discarded_messages, estimated_prompt_tokens = (
             multi_modal_truncate_prompt(
+                request=request,
                 messages=multi_modal_messages,
                 max_prompt_tokens=max_prompt_tokens,
-                initial_prompt_tokens=tokenizer.TOKENS_PER_REQUEST,
                 tokenizer=tokenizer,
             )
         )

diff --git a/aidial_adapter_openai/utils/tokenizer.py b/aidial_adapter_openai/utils/tokenizer.py
@@ -88,9 +88,7 @@ def _tokens_per_request_message_name(self) -> int:
         return 1
 
     def tokenize_request(
-        self,
-        original_request: dict,
-        messages: List[MessageType],
+        self, original_request: dict, messages: List[MessageType]
     ) -> int:
         tokens = self.TOKENS_PER_REQUEST
 

diff --git a/tests/test_discard_messages.py b/tests/test_discard_messages.py
@@ -141,7 +141,7 @@ def test_discarded_messages_without_error(
 ):
     tokenizer = PlainTextTokenizer(model="gpt-4")
     truncated_messages, discarded_messages, _used_tokens = (
-        plain_text_truncate_prompt(messages, max_prompt_tokens, tokenizer)
+        plain_text_truncate_prompt({}, messages, max_prompt_tokens, tokenizer)
     )
     assert (truncated_messages, discarded_messages) == response
 
@@ -157,5 +157,5 @@ def test_discarded_messages_with_error(
     tokenizer = PlainTextTokenizer(model="gpt-4")
 
     with pytest.raises(DialException) as e_info:
-        plain_text_truncate_prompt(messages, max_prompt_tokens, tokenizer)
+        plain_text_truncate_prompt({}, messages, max_prompt_tokens, tokenizer)
     assert e_info.value.message == error_message
diff --git a/tests/test_multimodal_truncate.py b/tests/test_multimodal_truncate.py
@@ -45,7 +45,7 @@ def test_multimodal_truncate_with_system_and_last_user_error():
         ),
     ]
     with pytest.raises(TruncatePromptSystemAndLastUserError):
-        multi_modal_truncate_prompt(transformations, 15, 0, tokenizer)
+        multi_modal_truncate_prompt({}, transformations, 15, tokenizer)
 
 
 def test_multimodal_truncate_with_system_error():
@@ -57,7 +57,7 @@ def test_multimodal_truncate_with_system_error():
         ),
     ]
     with pytest.raises(TruncatePromptSystemError):
-        multi_modal_truncate_prompt(transformations, 9, 3, tokenizer)
+        multi_modal_truncate_prompt({}, transformations, 9, tokenizer)
 
 
 @pytest.mark.parametrize(
@@ -194,9 +194,9 @@ def test_multimodal_truncate(
 ):
     truncated, actual_discarded_messages, actual_used_tokens = (
         multi_modal_truncate_prompt(
+            {},
             transformations,
             max_prompt_tokens,
-            initial_prompt_tokens=3,
             tokenizer=tokenizer,
         )
     )