Spaces:
Runtime error
Runtime error
from bot.session_manager import Session | |
from common.log import logger | |
class OpenAISession(Session): | |
def __init__(self, session_id, system_prompt=None, model= "text-davinci-003"): | |
super().__init__(session_id, system_prompt) | |
self.model = model | |
self.reset() | |
def __str__(self): | |
# 构造对话模型的输入 | |
''' | |
e.g. Q: xxx | |
A: xxx | |
Q: xxx | |
''' | |
prompt = "" | |
for item in self.messages: | |
if item['role'] == 'system': | |
prompt += item['content'] + "<|endoftext|>\n\n\n" | |
elif item['role'] == 'user': | |
prompt += "Q: " + item['content'] + "\n" | |
elif item['role'] == 'assistant': | |
prompt += "\n\nA: " + item['content'] + "<|endoftext|>\n" | |
if len(self.messages) > 0 and self.messages[-1]['role'] == 'user': | |
prompt += "A: " | |
return prompt | |
def discard_exceeding(self, max_tokens, cur_tokens= None): | |
precise = True | |
try: | |
cur_tokens = num_tokens_from_string(str(self), self.model) | |
except Exception as e: | |
precise = False | |
if cur_tokens is None: | |
raise e | |
logger.debug("Exception when counting tokens precisely for query: {}".format(e)) | |
while cur_tokens > max_tokens: | |
if len(self.messages) > 1: | |
self.messages.pop(0) | |
elif len(self.messages) == 1 and self.messages[0]["role"] == "assistant": | |
self.messages.pop(0) | |
if precise: | |
cur_tokens = num_tokens_from_string(str(self), self.model) | |
else: | |
cur_tokens = len(str(self)) | |
break | |
elif len(self.messages) == 1 and self.messages[0]["role"] == "user": | |
logger.warn("user question exceed max_tokens. total_tokens={}".format(cur_tokens)) | |
break | |
else: | |
logger.debug("max_tokens={}, total_tokens={}, len(conversation)={}".format(max_tokens, cur_tokens, len(self.messages))) | |
break | |
if precise: | |
cur_tokens = num_tokens_from_string(str(self), self.model) | |
else: | |
cur_tokens = len(str(self)) | |
return cur_tokens | |
# refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb | |
def num_tokens_from_string(string: str, model: str) -> int: | |
"""Returns the number of tokens in a text string.""" | |
import tiktoken | |
encoding = tiktoken.encoding_for_model(model) | |
num_tokens = len(encoding.encode(string,disallowed_special=())) | |
return num_tokens |