Jelajahi Sumber

feat: update counting tokens for new models

lanvent 2 tahun lalu
induk
melakukan
bb737a71d5
2 mengubah file dengan 9 tambahan dan 8 penghapusan
  1. 1 0
      bot/chatgpt/chat_gpt_bot.py
  2. 8 8
      bot/chatgpt/chat_gpt_session.py

+ 1 - 0
bot/chatgpt/chat_gpt_bot.py

@@ -121,6 +121,7 @@ class ChatGPTBot(Bot, OpenAIImage):
             if args is None:
                 args = self.args
             response = openai.ChatCompletion.create(api_key=api_key, messages=session.messages, **args)
+            # logger.debug("[CHATGPT] response={}".format(response))
             # logger.info("[ChatGPT] reply={}, total_tokens={}".format(response.choices[0]['message']['content'], response["usage"]["total_tokens"]))
             return {
                 "total_tokens": response["usage"]["total_tokens"],

+ 8 - 8
bot/chatgpt/chat_gpt_session.py

@@ -57,25 +57,25 @@ def num_tokens_from_messages(messages, model):
     """Returns the number of tokens used by a list of messages."""
     import tiktoken
 
-    if model == "gpt-3.5-turbo" or model == "gpt-35-turbo":
-        return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
-    elif model == "gpt-4":
-        return num_tokens_from_messages(messages, model="gpt-4-0314")
+    if model in ["gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613", "gpt-35-turbo", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613"]:
+        return num_tokens_from_messages(messages, model="gpt-3.5-turbo")
+    elif model in ["gpt-4-0314", "gpt-4-0613", "gpt-4-32k", "gpt-4-32k-0613"]:
+        return num_tokens_from_messages(messages, model="gpt-4")
 
     try:
         encoding = tiktoken.encoding_for_model(model)
     except KeyError:
         logger.debug("Warning: model not found. Using cl100k_base encoding.")
         encoding = tiktoken.get_encoding("cl100k_base")
-    if model == "gpt-3.5-turbo-0301":
+    if model == "gpt-3.5-turbo":
         tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
         tokens_per_name = -1  # if there's a name, the role is omitted
-    elif model == "gpt-4-0314":
+    elif model == "gpt-4":
         tokens_per_message = 3
         tokens_per_name = 1
     else:
-        logger.warn(f"num_tokens_from_messages() is not implemented for model {model}. Returning num tokens assuming gpt-3.5-turbo-0301.")
-        return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
+        logger.warn(f"num_tokens_from_messages() is not implemented for model {model}. Returning num tokens assuming gpt-3.5-turbo.")
+        return num_tokens_from_messages(messages, model="gpt-3.5-turbo")
     num_tokens = 0
     for message in messages:
         num_tokens += tokens_per_message