Spaces:
Runtime error
Runtime error
| import time | |
| import json | |
| import yaml | |
| from typing import Union | |
| import os | |
| from transformers import pipeline | |
| class ChatBot(object): | |
| def __init__(self): | |
| self.model = pipe = pipeline("text-generation", model="Team-Coffee-Gym/CoffeeGym") | |
| self.chat_history = [] | |
| self.history = self._set_initial_history() | |
| def _set_initial_history(self): | |
| return ["You are an exceptionally intelligent coding assistant developed by DLI lab that consistently delivers accurate and reliable responses to user instructions. If somebody asks you who are you, answer as 'AI programming assistant based on DLI Lab'.\n\n"] | |
| def set_model_input(self, input_text = None): | |
| model_input = [] | |
| if input_text is not None: | |
| self.history.append(input_text) | |
| model_input.append({ | |
| "role": "system", | |
| "content": self.history[0] | |
| }) | |
| chat_history = self.history[1:] | |
| for i in range(len(chat_history)): | |
| if i % 2 == 0: | |
| model_input.append({ | |
| "role": "user", | |
| "content": chat_history[i] | |
| }) | |
| else: | |
| model_input.append({ | |
| "role": "assistant", | |
| "content": chat_history[i] | |
| }) | |
| return model_input | |
| def chat(self, chat_history, input_text): | |
| self.chat_history = chat_history | |
| model_input = self.set_model_input(input_text) | |
| response = self.model(model_input) | |
| if response is not None: | |
| self.history.append(response) | |
| self.chat_history = self.chat_history + [(input_text, response)] | |
| self.log_chat() | |
| return self.chat_history | |
| def regenerate(self, chat_history, input_text): | |
| self.chat_history = chat_history[:-1] | |
| self.history = self.history[:-2] | |
| model_input = self.set_model_input(None) | |
| response = self.model.invoke(model_input) | |
| if response is not None: | |
| self.history.append(response) | |
| self.chat_history = self.chat_history + [(input_text, response)] | |
| return self.chat_history | |
| def clear_chat(self): | |
| self.chat_history = [] | |
| self.history = self._set_initial_history() | |
| return self.chat_history |