Spaces:
Sleeping
Sleeping
| from bot.session_manager import Session | |
| from common.log import logger | |
| ''' | |
| e.g. [ | |
| {"role": "system", "content": "You are a helpful assistant."}, | |
| {"role": "user", "content": "Who won the world series in 2020?"}, | |
| {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}, | |
| {"role": "user", "content": "Where was it played?"} | |
| ] | |
| ''' | |
| class ChatGPTSession(Session): | |
| def __init__(self, session_id, system_prompt=None, model= "gpt-3.5-turbo"): | |
| super().__init__(session_id, system_prompt) | |
| self.model = model | |
| self.reset() | |
| def discard_exceeding(self, max_tokens, cur_tokens= None): | |
| precise = True | |
| try: | |
| cur_tokens = num_tokens_from_messages(self.messages, 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) > 2: | |
| self.messages.pop(1) | |
| 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) | |
| else: | |
| cur_tokens = cur_tokens - max_tokens | |
| break | |
| elif len(self.messages) == 2 and self.messages[1]["role"] == "user": | |
| logger.warn("user message exceed max_tokens. total_tokens={}".format(cur_tokens)) | |
| break | |
| else: | |
| 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) | |
| else: | |
| cur_tokens = cur_tokens - max_tokens | |
| 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_messages(messages, model): | |
| """Returns the number of tokens used by a list of messages.""" | |
| import tiktoken | |
| 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": | |
| 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") | |
| elif model == "gpt-3.5-turbo-0301": | |
| 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": | |
| 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") | |
| num_tokens = 0 | |
| for message in messages: | |
| num_tokens += tokens_per_message | |
| for key, value in message.items(): | |
| num_tokens += len(encoding.encode(value)) | |
| if key == "name": | |
| num_tokens += tokens_per_name | |
| num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> | |
| return num_tokens |