選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。

chat_gpt_session.py 3.7KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. from bot.session_manager import Session
  2. from common.log import logger
  3. """
  4. e.g. [
  5. {"role": "system", "content": "You are a helpful assistant."},
  6. {"role": "user", "content": "Who won the world series in 2020?"},
  7. {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
  8. {"role": "user", "content": "Where was it played?"}
  9. ]
  10. """
  11. class ChatGPTSession(Session):
  12. def __init__(self, session_id, system_prompt=None, model="gpt-3.5-turbo"):
  13. super().__init__(session_id, system_prompt)
  14. self.model = model
  15. self.reset()
  16. def discard_exceeding(self, max_tokens, cur_tokens=None):
  17. precise = True
  18. try:
  19. cur_tokens = self.calc_tokens()
  20. except Exception as e:
  21. precise = False
  22. if cur_tokens is None:
  23. raise e
  24. logger.debug(
  25. "Exception when counting tokens precisely for query: {}".format(e)
  26. )
  27. while cur_tokens > max_tokens:
  28. if len(self.messages) > 2:
  29. self.messages.pop(1)
  30. elif len(self.messages) == 2 and self.messages[1]["role"] == "assistant":
  31. self.messages.pop(1)
  32. if precise:
  33. cur_tokens = self.calc_tokens()
  34. else:
  35. cur_tokens = cur_tokens - max_tokens
  36. break
  37. elif len(self.messages) == 2 and self.messages[1]["role"] == "user":
  38. logger.warn(
  39. "user message exceed max_tokens. total_tokens={}".format(cur_tokens)
  40. )
  41. break
  42. else:
  43. logger.debug(
  44. "max_tokens={}, total_tokens={}, len(messages)={}".format(
  45. max_tokens, cur_tokens, len(self.messages)
  46. )
  47. )
  48. break
  49. if precise:
  50. cur_tokens = self.calc_tokens()
  51. else:
  52. cur_tokens = cur_tokens - max_tokens
  53. return cur_tokens
  54. def calc_tokens(self):
  55. return num_tokens_from_messages(self.messages, self.model)
  56. # refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
  57. def num_tokens_from_messages(messages, model):
  58. """Returns the number of tokens used by a list of messages."""
  59. import tiktoken
  60. try:
  61. encoding = tiktoken.encoding_for_model(model)
  62. except KeyError:
  63. logger.debug("Warning: model not found. Using cl100k_base encoding.")
  64. encoding = tiktoken.get_encoding("cl100k_base")
  65. if model == "gpt-3.5-turbo" or model == "gpt-35-turbo":
  66. return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
  67. elif model == "gpt-4":
  68. return num_tokens_from_messages(messages, model="gpt-4-0314")
  69. elif model == "gpt-3.5-turbo-0301":
  70. tokens_per_message = (
  71. 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
  72. )
  73. tokens_per_name = -1 # if there's a name, the role is omitted
  74. elif model == "gpt-4-0314":
  75. tokens_per_message = 3
  76. tokens_per_name = 1
  77. else:
  78. logger.warn(
  79. f"num_tokens_from_messages() is not implemented for model {model}. Returning num tokens assuming gpt-3.5-turbo-0301."
  80. )
  81. return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
  82. num_tokens = 0
  83. for message in messages:
  84. num_tokens += tokens_per_message
  85. for key, value in message.items():
  86. num_tokens += len(encoding.encode(value))
  87. if key == "name":
  88. num_tokens += tokens_per_name
  89. num_tokens += 3 # every reply is primed with <|start|>assistant<|message|>
  90. return num_tokens