Merge pull request #2339 from 6vision/master

Optimize Gemini model character statistics
zhayujie · Oct 14, 2024 · fcc4d02 · fcc4d02
2 parents 883f0d4 + f4a5f00
commit fcc4d02
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bot/chatgpt/chat_gpt_session.py b/bot/chatgpt/chat_gpt_session.py
@@ -57,7 +57,7 @@ def calc_tokens(self):
 def num_tokens_from_messages(messages, model):
     """Returns the number of tokens used by a list of messages."""
 
-    if model in ["wenxin", "xunfei", const.GEMINI]:
+    if model in ["wenxin", "xunfei"] or model.startswith(const.GEMINI):
         return num_tokens_by_character(messages)
 
     import tiktoken