From 051dc133228f5b64355e0c323ec45c8bf605c19c Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <80892890+shaohuzhang1@users.noreply.github.com>
Date: Thu, 13 Feb 2025 15:40:09 +0800
Subject: [PATCH] perf: Add custom token_ids method (#2264)

---
 apps/setting/models_provider/impl/base_chat_open_ai.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/apps/setting/models_provider/impl/base_chat_open_ai.py b/apps/setting/models_provider/impl/base_chat_open_ai.py
index 0989517a75d..d00002ed749 100644
--- a/apps/setting/models_provider/impl/base_chat_open_ai.py
+++ b/apps/setting/models_provider/impl/base_chat_open_ai.py
@@ -15,8 +15,14 @@
 from common.config.tokenizer_manage_config import TokenizerManage
 
 
+def custom_get_token_ids(text: str):
+    tokenizer = TokenizerManage.get_tokenizer()
+    return tokenizer.encode(text)
+
+
 class BaseChatOpenAI(ChatOpenAI):
     usage_metadata: dict = {}
+    custom_get_token_ids = custom_get_token_ids
 
     def get_last_generation_info(self) -> Optional[Dict[str, Any]]:
         return self.usage_metadata