Update app.py

app.py

@@ -1,17 +1,13 @@
-#importing libraries
 import gradio as gr
 import tensorflow.keras as keras
 import time
 import keras_nlp
 import os
 
-
 model_path = "Zul001/HydroSense_Gemma_Finetuned_Model"
 gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(f"hf://{model_path}")
 
-
 reset_triggered = False
-
 custom_css = """
 @import url('https://fonts.googleapis.com/css2?family=Edu+AU+VIC+WA+NT+Dots:[email protected]&family=Give+You+Glory&family=Sofia&family=Sunshiney&family=Vujahday+Script&display=swap');
 .gradio-container, .gradio-container * {
@@ -31,27 +27,53 @@ function refresh() {
 }
 """
 
-
 previous_sessions = []
 
+class ChatState():
+    __START_TURN_USER__ = "<start_of_turn>user\n"
+    __START_TURN_MODEL__ = "<start_of_turn>model\n"
+    __END_TURN__ = "<end_of_turn>\n"
+
+    def __init__(self, model, system=""):
+        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__)
+
+    def add_to_history_as_model(self, message):
+        self.history.append(self.__START_TURN_MODEL__ + message)
+
+    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.system + "\n" + prompt
+        return prompt
+
+    def send_message(self, message):
+        self.add_to_history_as_user(message)
+        prompt = self.get_full_prompt()
+        response = self.model.generate(prompt, max_length=2048)
+        result = response.replace(prompt, "")  # Extract only the new response
+        self.add_to_history_as_model(result)
+        return result
+
 def post_process_output(prompt, result):
-    # Remove the prompt if it's repeated at the beginning of the answer
     answer = result.strip()
     if answer.startswith(prompt):
         answer = answer[len(prompt):].strip()
 
-    # Remove any leading colons or whitespace
     answer = answer.lstrip(':')
-
-    # Ensure the answer starts with a capital letter
     answer = answer.capitalize()
 
-    # Ensure the answer ends with a period if it doesn't already
     if not answer.endswith('.'):
         answer += '.'
 
     return f"{answer}"
-    
 
 def add_session(prompt):
     global previous_sessions
@@ -62,60 +84,36 @@ def add_session(prompt):
         
     return "\n".join(previous_sessions)  # Return only the session logs as a string
 
+def clear_sessions():
+    global previous_sessions
+    previous_sessions.clear()
+    return "\n".join(previous_sessions)
+
+def clear_fields():
+    global reset_triggered
+    reset_triggered = True
+    return "", ""  # Return empty strings to clear the prompt and output fields
 
+# Initialize the ChatState
+chat_state = ChatState(gemma_lm)
 
 def inference(prompt):
     global reset_triggered
     if reset_triggered:
-        #do nothing
         return "", ""
     
-    prompt_text = prompt
-    generated_text = gemma_lm.generate(prompt_text)
-    
-    #Apply post-processing
-    formatted_output = post_process_output(prompt_text, generated_text)
-    print(formatted_output)
+    chat_state.send_message(prompt)  # Process the user's message
+
+    # Post-process the output from the model
+    formatted_output = post_process_output(chat_state.get_full_prompt(), chat_state.get_history())
     
-    #adding a bit of delay
+    # Apply a bit of delay for a realistic response time
     time.sleep(1)
+    
     result = formatted_output
-    sessions = add_session(prompt_text)
+    sessions = add_session(chat_state.get_history())
     return result, sessions
 
-
-# def inference(prompt):
-    
-#     time.sleep(1)
-#     result = "Your Result"
-#     # sessions = add_session(prompt)
-#     return result
-
-     
-# def remember(prompt, result):
-#     global memory
-#     # Store the session as a dictionary
-#     session = {'prompt': prompt, 'result': result}
-#     memory.append(session)
-
-#     # Update previous_sessions for display
-#     session_display = [f"Q: {s['prompt']} \nA: {s['result']}" for s in memory]
-    
-#     return "\n\n".join(session_display)  # Return formatted sessions as a string
-
-
-
-def clear_sessions():
-    global previous_sessions
-    previous_sessions.clear()
-    return "\n".join(previous_sessions)
-
-def clear_fields():
-    global reset_triggered
-    reset_triggered = True
-    return "", ""  # Return empty strings to clear the prompt and output fields
-
-
 with gr.Blocks(theme='gradio/soft', css=custom_css) as demo:
     
     gr.Markdown("<center><h1>HydroSense LLM Demo</h1></center>")
@@ -135,9 +133,6 @@ with gr.Blocks(theme='gradio/soft', css=custom_css) as demo:
                 generate_btn = gr.Button("Generate Answer", variant="primary", size="sm")
                 reset_btn = gr.Button("Clear Content", variant="secondary", size="sm", elem_id="primary")
 
-        
-
-
     generate_btn.click(
         fn=inference,
         inputs=[prompt],
@@ -156,19 +151,16 @@ with gr.Blocks(theme='gradio/soft', css=custom_css) as demo:
         outputs=[prompt, output]
     )
 
-
-    # Button to clear the prompt and output fields
     add_button.click(
         fn=clear_fields,  # Only call the clear_fields function
         inputs=None,      # No inputs needed
         outputs=[prompt, output]  # Clear the prompt and output fields
 )
 
-
     clear_session.click(
         fn=clear_sessions,
         inputs=None,
         outputs=[session_list]
     )
 
-demo.launch(share=True)
+demo.launch(share=True)