Update app.py

app.py

@@ -1,13 +1,17 @@
+#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 * {
@@ -27,64 +31,27 @@ function refresh() {
 }
 """
 
-previous_sessions = []
-
-class ChatState:
-    def __init__(self, model, tokenizer, system=""):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.system = system
-        self.history = []
-
-    def add_to_history(self, role, message):
-        self.history.append({"role": role, "content": message})
-
-    def get_full_prompt(self):
-        prompt = ""
-        if self.system:
-            prompt += f"System: {self.system}\n\n"
-        for message in self.history:
-            prompt += f"{message['role'].capitalize()}: {message['content']}\n"
-        prompt += "Model: "
-        return prompt
-
-    def send_message(self, message):
-        self.add_to_history("user", message)
-        prompt = self.get_full_prompt()
-        
-        inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-        
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=512,
-                num_return_sequences=1,
-                do_sample=True,
-                temperature=0.7
-            )
-        
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response = response.replace(prompt, "").strip()
-        
-        self.add_to_history("model", response)
-        return response
-
-# Initialize the ChatState
-chat_state = ChatState(model, tokenizer, system="You are a helpful AI assistant.")
 
+previous_sessions = []
 
 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
@@ -95,33 +62,60 @@ 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 "", ""
     
-    result = chat_state.send_message(prompt)
+    prompt_text = prompt
+    generated_text = gemma_lm.generate(prompt_text)
     
-    # Apply a bit of delay for a realistic response time
-    time.sleep(1)
+    #Apply post-processing
+    formatted_output = post_process_output(prompt_text, generated_text)
+    print(formatted_output)
     
-    sessions = add_session(prompt)
+    #adding a bit of delay
+    time.sleep(1)
+    result = formatted_output
+    sessions = add_session(prompt_text)
     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>")
@@ -141,6 +135,9 @@ 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],
@@ -159,12 +156,15 @@ 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,