| class ChatState(): |
| """ |
| Manages the conversation history for a turn-based chatbot |
| Follows the turn-based conversation guidelines for the Gemma family of models |
| documented at https://ai.google.dev/gemma/docs/formatting |
| """ |
|
|
|
|
| __START_TURN_USER__ = "Instruction:\n" |
| __START_TURN_MODEL__ = "\n\nResponse:\n" |
| __END_TURN__ = "" |
|
|
|
|
| def __init__(self, model, system=""): |
| """ |
| Initializes the chat state. |
| |
| Args: |
| model: The language model to use for generating responses. |
| system: (Optional) System instructions or bot description. |
| """ |
| self.model = model |
| self.system = system |
| self.history = [] |
|
|
| def add_to_history_as_user(self, message): |
| """ |
| Adds a user message to the history with start/end turn markers. |
| """ |
| self.history.append(self.__START_TURN_USER__ + message + self.__END_TURN__) |
|
|
| def add_to_history_as_model(self, message): |
| """ |
| Adds a model response to the history with the start turn marker. |
| Model will generate end turn marker. |
| """ |
| self.history.append(self.__START_TURN_MODEL__ + message+ "\n") |
|
|
| def get_history(self): |
| """ |
| Returns the entire chat history as a single string. |
| """ |
| return "".join([*self.history]) |
|
|
| def get_full_prompt(self): |
| """ |
| Builds the prompt for the language model, including history and system description. |
| """ |
| prompt = self.get_history() + self.__START_TURN_MODEL__ |
| if len(self.system)>0: |
| prompt = self.system + "\n" + 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=4096) |
| result = response.replace(prompt, "") |
| self.add_to_history_as_model(result) |
| return result |
|
|
