Spaces:
Running
on
TPU v5e
Running
on
TPU v5e
| # chat helper | |
| class ChatState: | |
| def __init__(self, model, system="", chat_template="auto"): | |
| chat_template = ( | |
| type(model).__name__ if chat_template == "auto" else chat_template | |
| ) | |
| if chat_template == "Llama3CausalLM": | |
| self.__START_TURN_SYSTEM__ = ( | |
| "<|start_header_id|>system<|end_header_id|>\n\n" | |
| ) | |
| self.__START_TURN_USER__ = ( | |
| "<|start_header_id|>user<|end_header_id|>\n\n" | |
| ) | |
| self.__START_TURN_MODEL__ = ( | |
| "<|start_header_id|>assistant<|end_header_id|>\n\n" | |
| ) | |
| self.__END_TURN_SYSTEM__ = "<|eot_id|>" | |
| self.__END_TURN_USER__ = "<|eot_id|>" | |
| self.__END_TURN_MODEL__ = "<|eot_id|>" | |
| print("Using chat template for: Llama") | |
| elif chat_template == "GemmaCausalLM": | |
| self.__START_TURN_SYSTEM__ = "" | |
| self.__START_TURN_USER__ = "<start_of_turn>user\n" | |
| self.__START_TURN_MODEL__ = "<start_of_turn>model\n" | |
| self.__END_TURN_SYSTEM__ = "\n" | |
| self.__END_TURN_USER__ = "<end_of_turn>\n" | |
| self.__END_TURN_MODEL__ = "<end_of_turn>\n" | |
| print("Using chat template for: Gemma") | |
| elif chat_template == "MistralCausalLM": | |
| self.__START_TURN_SYSTEM__ = "" | |
| self.__START_TURN_USER__ = "[INST]" | |
| self.__START_TURN_MODEL__ = "" | |
| self.__END_TURN_SYSTEM__ = "<s>" | |
| self.__END_TURN_USER__ = "[/INST]" | |
| self.__END_TURN_MODEL__ = "</s>" | |
| print("Using chat template for: Mistral") | |
| elif chat_template == "Vicuna": | |
| self.__START_TURN_SYSTEM__ = "" | |
| self.__START_TURN_USER__ = "USER: " | |
| self.__START_TURN_MODEL__ = "ASSISTANT: " | |
| self.__END_TURN_SYSTEM__ = "\n\n" | |
| self.__END_TURN_USER__ = "\n" | |
| self.__END_TURN_MODEL__ = "</s>\n" | |
| print("Using chat template for : Vicuna") | |
| else: | |
| assert (0, "Unknown turn tags for this model class") | |
| self.model = model | |
| self.system = system | |
| self.history = [] | |
| def add_to_history_as_user(self, message): | |
| self.history.append( | |
| self.__START_TURN_USER__ + message + self.__END_TURN_USER__ | |
| ) | |
| def add_to_history_as_model(self, message): | |
| self.history.append( | |
| self.__START_TURN_MODEL__ + message + self.__END_TURN_MODEL__ | |
| ) | |
| def get_history(self): | |
| return "".join([*self.history]) | |
| def get_full_prompt(self): | |
| prompt = self.get_history() + self.__START_TURN_MODEL__ | |
| if len(self.system) > 0: | |
| prompt = ( | |
| self.__START_TURN_SYSTEM__ | |
| + self.system | |
| + self.__END_TURN_SYSTEM__ | |
| + prompt | |
| ) | |
| return prompt | |
| def send_message(self, message): | |
| """ | |
| Handles sending a user message and getting a model response. | |
| Args: | |
| message: The user's message. | |
| Returns: | |
| The model's response. | |
| """ | |
| self.add_to_history_as_user(message) | |
| prompt = self.get_full_prompt() | |
| response = self.model.generate( | |
| prompt, max_length=2048, strip_prompt=True | |
| ) | |
| self.add_to_history_as_model(response) | |
| return (message, response) | |