12345678910111213141516171819202122232425 |
- class Chat:
- """
- Handles the chat using to avoid OOM issues.
- """
- def __init__(self, size):
- self.size = size
- self.init_chat_message = None
- # maxlen is necessary pair, since a each new step we add an prompt and assitant answer
- self.buffer = []
- def append(self, item):
- self.buffer.append(item)
- if len(self.buffer) == 2 * (self.size + 1):
- self.buffer.pop(0)
- self.buffer.pop(0)
- def init_chat(self, init_chat_message):
- self.init_chat_message = init_chat_message
- def to_list(self):
- if self.init_chat_message:
- return [self.init_chat_message] + self.buffer
- else:
- return self.buffer
|