kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-20 21:57:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: add calc_tokens method on session

Browse Source
This commit is contained in:
lanvent
2023-04-13 20:06:33 +08:00
parent 9cef559a05
commit 2989249e4b
5 changed files with 31 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
bot/chatgpt/chat_gpt_session.py
View File

@@ -17,7 +17,7 @@ class ChatGPTSession(Session):
def discard_exceeding(self, max_tokens, cur_tokens= None):
precise = True
try:
cur_tokens = num_tokens_from_messages(self.messages, self.model)
cur_tokens = self.calc_tokens()
except Exception as e:
precise = False
if cur_tokens is None:
@@ -29,7 +29,7 @@ class ChatGPTSession(Session):
elif len(self.messages) == 2 and self.messages[1]["role"] == "assistant":
self.messages.pop(1)
if precise:
cur_tokens = num_tokens_from_messages(self.messages, self.model)
cur_tokens = self.calc_tokens()
else:
cur_tokens = cur_tokens - max_tokens
break
@@ -40,11 +40,14 @@ class ChatGPTSession(Session):
logger.debug("max_tokens={}, total_tokens={}, len(messages)={}".format(max_tokens, cur_tokens, len(self.messages)))
break
if precise:
cur_tokens = num_tokens_from_messages(self.messages, self.model)
cur_tokens = self.calc_tokens()
else:
cur_tokens = cur_tokens - max_tokens
return cur_tokens
def calc_tokens(self):
return num_tokens_from_messages(self.messages, self.model)
# refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
def num_tokens_from_messages(messages, model):