perf: Add custom token_ids method (#2264)

1Panel-dev · Feb 13, 2025 · 051dc13 · 051dc13
1 parent dac0a9c
commit 051dc13
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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