Update app.py

app.py

@@ -58,9 +58,7 @@ def log_feedback(bias_mode, prompt, response, thumb):
     if not prompt or not response:
         return
     with open(FEEDBACK_FILE, "a", newline="", encoding="utf-8") as f:
-        writer = csv.writer(
-            f
-        )
+        writer = csv.writer(f)
         writer.writerow(
             [
                 datetime.utcnow().isoformat(),
@@ -118,6 +116,10 @@ def build_context(messages, user_message, bias_mode):
 
 
 def generate_response(user_message, messages, bias_mode):
+    """
+    - messages: list of message dicts (Chatbot "messages" format)
+    Returns: (cleared textbox, updated messages, last_user, last_bot)
+    """
     if not user_message.strip():
         return "", messages, messages, "", ""
 
@@ -133,12 +135,14 @@ def generate_response(user_message, messages, bias_mode):
     )
 
     full_text = outputs[0]["generated_text"]
+
+    # ✅ Use the *last* Assistant: block (the new reply)
     if "Assistant:" in full_text:
-        bot_part = full_text.split("Assistant:", 1)[1]
+        bot_part = full_text.rsplit("Assistant:", 1)[1]
     else:
         bot_part = full_text
 
-    # *** NEW: cut off if model starts writing a new "User:" line ***
+    # ✅ Cut off if the model starts writing a new "User:" line
     bot_part = bot_part.split("\nUser:")[0].strip()
 
     bot_reply = bot_part
@@ -152,7 +156,6 @@ def generate_response(user_message, messages, bias_mode):
     return "", messages, messages, user_message, bot_reply
 
 
-
 def handle_thumb(thumb_value, last_user, last_bot, bias_mode):
     """
     Called when user clicks 👍 or 👎.
@@ -282,18 +285,20 @@ def run_bias_probe(bias_mode: str) -> str:
             prompt_text,
             max_new_tokens=120,
             do_sample=True,
-            top_p=0.95,
-            temperature=0.8,
+            top_p=0.9,
+            temperature=0.7,
             pad_token_id=tokenizer.eos_token_id,
         )
 
         full_text = outputs[0]["generated_text"]
         if "Assistant:" in full_text:
-            answer = full_text.split("Assistant:")[-1].strip()
+            answer_part = full_text.rsplit("Assistant:", 1)[1]
         else:
-            answer = full_text.strip()
+            answer_part = full_text
+
+        answer_part = answer_part.split("\nUser:")[0].strip()
 
-        reports.append(f"**Q:** {q}\n\n**A:** {answer}\n")
+        reports.append(f"**Q:** {q}\n\n**A:** {answer_part}\n")
 
     header = f"### Bias probe results (mode: *{bias_mode}*)\n"
     return header + "\n---\n".join(reports)
@@ -333,7 +338,6 @@ with gr.Blocks() as demo:
             label="Current bias target",
         )
 
-    # Chatbot now uses default "messages" format (list of dicts with role/content)
     chatbot = gr.Chatbot(height=400, label="EnergyBiasShifter")
 
     msg = gr.Textbox(