chat.py 724 B

12345678910111213141516171819202122232425
  1. class Chat:
  2. """
  3. Handles the chat using to avoid OOM issues.
  4. """
  5. def __init__(self, size):
  6. self.size = size
  7. self.init_chat_message = None
  8. # maxlen is necessary pair, since a each new step we add an prompt and assitant answer
  9. self.buffer = []
  10. def append(self, item):
  11. self.buffer.append(item)
  12. if len(self.buffer) == 2 * (self.size + 1):
  13. self.buffer.pop(0)
  14. self.buffer.pop(0)
  15. def init_chat(self, init_chat_message):
  16. self.init_chat_message = init_chat_message
  17. def to_list(self):
  18. if self.init_chat_message:
  19. return [self.init_chat_message] + self.buffer
  20. else:
  21. return self.buffer