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