Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 os.environ["KERAS_BACKEND"] = "tensorflow"                                                                           
 import keras
 import keras_nlp
-
+import utils
 
 css = """
 html, body {
@@ -35,18 +35,24 @@ body::before {
 }
 """
 
-keras.config.set_floatx("bfloat16")
+#keras.config.set_floatx("bfloat16") uncomment to run at half precision.
 gemma_lm = keras_nlp.models.CausalLM.from_preset("hf://sultan-hassan/CosmoGemma_2b_en")
+chat = utils.ChatState(gemma_lm) # adding a chat helper to manage the conversation state (modified version from https://github.com/google-gemini/gemma-cookbook/blob/main/Gemma/Keras_Gemma_2_Quickstart_Chat.ipynb)
+
+def launch(message):
 
-def launch(input):
-    template = "Instruction:\n{instruction}\n\nResponse:\n{response}"
-    prompt = template.format(
-        instruction=input,                                                                   
-        response="",
-    )
-    out = gemma_lm.generate(prompt, max_length=1024)
-    ind = out.index('Response') + len('Response')+2
-    return out[ind:]
+    # Uncomment for QA system without chat history/memory
+    #template = "Instruction:\n{instruction}\n\nResponse:\n{response}"
+    #prompt = template.format(
+    #    instruction=message,                                                                   
+    #    response="",
+    #)
+    #out = gemma_lm.generate(prompt, max_length=1024)
+    #ind = out.index('\n\nResponse:\n') + len('\n\nResponse:\n')
+    #return out[ind:]
+    
+    # This is to convert QA system to a chatbot
+    return chat.send_message(message)
 
 iface = gr.Interface(launch,
                      inputs="text",