|
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667 |
- 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
|
