Final touches

app.py

@@ -1,4 +1,6 @@
 import os
+import json
+import random
 from threading import Thread
 
 
@@ -51,10 +53,10 @@ class HfModelWrapper:
             "inputs": input_ids,
             "streamer": streamer,
             "pad_token_id": self.tokenizer.eos_token_id,
-            "max_length": 8192,
+            "max_length": 32768,
             "temperature": 0.1,
             "top_p": 0.8,
-            "repetition_penalty": 1.25,
+            "repetition_penalty": 1.1,
         }
 
         thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
@@ -486,20 +488,36 @@ paper_plane_svg = """<svg xmlns="http://www.w3.org/2000/svg" width="20" height="
 </svg>"""
 
 wrapper = HfModelWrapper()
+with open("./question_mapping.json", "r") as f:
+    database = json.load(f)
+database = {k.lower(): v for k, v in database.items()}
 
 
 # Mock data function for chatbot
 def send_message(message, history):
     #if not message:
     #    return history
-    history.append({"role": "user", "content": message})
     #history.append({"role": "assistant", "content": f"This is a response about: {message}"})
     #return history
+    
+    history.append({"role": "user", "content": message})
+    # Recover extra data
+    if message.lower() in database:
+        context = random.choice(database[message.lower()])
+    message += " \n"+context
     response_generator = wrapper.generate(message, history)
     partial = ""
+    idx = 0
     for t in response_generator:
-        partial += t
-        yield history + [{"role": "assistant", "content": partial}]    
+        if idx <= 5:
+            idx += 1
+            continue
+        else:
+            if t != "<|eot_id|>":
+                partial += t
+                yield history + [{"role": "assistant", "content": partial}]    
+            else:
+                continue
     
 
 # Create the dashboard
@@ -542,18 +560,36 @@ with gr.Blocks(theme=ArtemisiaTheme(), css=custom_css) as demo:
                 results_placeholder = gr.Markdown("Ask a question in the chat to analyze risk data and see insights here", elem_classes="placeholder-text")
 
     # Wire up the chat functionality
+    #send_button.click(
+    #    fn=send_message,
+    #    inputs=[chat_input, chatbot],
+    #    outputs=[chatbot]
+    #).then(lambda: "", None, chat_input)
+    
+    #chat_input.submit(
+    #    fn=send_message,
+    #    inputs=[chat_input, chatbot],
+    #    outputs=[chatbot]
+    #).then(lambda: "", None, chat_input)
     send_button.click(
+        lambda msg: (gr.update(value=""), msg),  # clears textbox immediately
+        inputs=[chat_input],
+        outputs=[chat_input, gr.State()],  # dummy state to capture message
+    ).then(
         fn=send_message,
         inputs=[chat_input, chatbot],
         outputs=[chatbot]
-    ).then(lambda: "", None, chat_input)
+    )
     
     chat_input.submit(
+        lambda msg: (gr.update(value=""), msg),
+        inputs=[chat_input],
+        outputs=[chat_input, gr.State()],
+    ).then(
         fn=send_message,
         inputs=[chat_input, chatbot],
         outputs=[chatbot]
-    ).then(lambda: "", None, chat_input)
-
+    )
     # JavaScript for UI enhancements
     gr.HTML("""
     <script>