|
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283 |
- 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 = self.calc_tokens()
- 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 = self.calc_tokens()
- 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 = self.calc_tokens()
- else:
- cur_tokens = len(str(self))
- return cur_tokens
-
- def calc_tokens(self):
- return num_tokens_from_string(str(self), self.model)
-
-
-
- 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
|