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