Update prompt formating

app.py

@@ -82,9 +82,14 @@ def generate_response(message, chat_history, epsilon):
     # Set model thresholds based on epsilon
     model.head_thresholds = torch.tensor(epsilon_thresholds[epsilon])
     
+    # Format the prompt with chat history
+    formatted_prompt = ""
+    for user_msg, assistant_msg in chat_history:
+        formatted_prompt += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
+    formatted_prompt += f"User: {message}\nAssistant:"
+    
     full_response = ""
-    chat_history = chat_history or []
-    inputs = tokenizer.encode(message, return_tensors="pt").to(device)
+    inputs = tokenizer.encode(formatted_prompt, return_tensors="pt").to(device)
     
     while not stop_generation:
         inputs = truncate_context(inputs)
@@ -115,7 +120,7 @@ def generate_response(message, chat_history, epsilon):
         
         new_history = chat_history + [(message, full_response)]
         yield new_history, new_history, gr.update(value=create_plot())
-
+        
 def stop_gen():
     global stop_generation
     stop_generation = True