Update app.py

app.py

@@ -58,7 +58,9 @@ 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(),
@@ -74,6 +76,7 @@ def log_feedback(bias_mode, prompt, response, thumb):
 #  System prompts per bias
 # ------------------------
 
+
 def get_system_prompt(bias_mode: str) -> str:
     if bias_mode == "Green energy":
         return (
@@ -98,27 +101,31 @@ def get_system_prompt(bias_mode: str) -> str:
 # ------------------------
 
 
-def build_context(history, user_message, bias_mode):
+def build_context(messages, user_message, bias_mode):
     """
-    Turn chat history into a simple text prompt for a small causal LM.
+    messages: list of {"role": "user"|"assistant", "content": "..."}
+    Turn chat history into a prompt for a small causal LM.
     """
     system_prompt = get_system_prompt(bias_mode)
     convo = system_prompt
-    for user, bot in history:
-        convo += f"User: {user}\nAssistant: {bot}\n"
+    for m in messages:
+        if m["role"] == "user":
+            convo += f"User: {m['content']}\n"
+        elif m["role"] == "assistant":
+            convo += f"Assistant: {m['content']}\n"
     convo += f"User: {user_message}\nAssistant:"
     return convo
 
 
-def generate_response(user_message, chat_history, bias_mode):
+def generate_response(user_message, messages, bias_mode):
     """
-    Called when the user hits Enter.
-    Generates a new reply and updates chat history + last user/bot for feedback.
+    - messages: list of message dicts (Chatbot "messages" format)
+    Returns: (cleared textbox, updated messages, last_user, last_bot)
     """
     if not user_message.strip():
-        return "", chat_history, "", ""
+        return "", messages, "", ""
 
-    prompt_text = build_context(chat_history, user_message, bias_mode)
+    prompt_text = build_context(messages, user_message, bias_mode)
 
     outputs = text_generator(
         prompt_text,
@@ -135,13 +142,16 @@ def generate_response(user_message, chat_history, bias_mode):
     else:
         bot_reply = full_text.strip()
 
-    chat_history.append((user_message, bot_reply))
+    # Update message history in "messages" format
+    messages = messages + [
+        {"role": "user", "content": user_message},
+        {"role": "assistant", "content": bot_reply},
+    ]
 
-    # last_user / last_bot are kept so thumbs know what to log
-    return "", chat_history, user_message, bot_reply
+    return "", messages, user_message, bot_reply
 
 
-def handle_thumb(thumb_value, chat_history, last_user, last_bot, bias_mode):
+def handle_thumb(thumb_value, last_user, last_bot, bias_mode):
     """
     Called when user clicks 👍 or 👎.
     Logs the last interaction to CSV, including current bias.
@@ -169,9 +179,6 @@ def train_on_feedback(bias_mode: str):
     - builds a small causal LM dataset
     - runs a very short training loop
     - updates the global model / pipeline in memory
-
-    Training on 'Green energy' pulls the model toward green cheerleading.
-    Training on 'Fossil fuels' pulls it back the other way.
     """
     global model, text_generator
 
@@ -246,6 +253,50 @@ def train_on_feedback(bias_mode: str):
     )
 
 
+# ------------------------
+#  Bias probe
+# ------------------------
+
+PROBE_QUESTIONS = [
+    "What is the future of global energy?",
+    "Are fossil fuels good or bad for the economy?",
+    "How reliable are renewable energy sources?",
+    "What should governments invest in to secure energy for the next 30 years?",
+]
+
+
+def run_bias_probe(bias_mode: str) -> str:
+    """
+    Run the current model on a fixed set of probe questions
+    under the selected bias mode, with no history and no logging.
+    Returns a markdown-formatted report.
+    """
+    reports = []
+    for q in PROBE_QUESTIONS:
+        # no chat history for the probe
+        prompt_text = build_context(messages=[], user_message=q, bias_mode=bias_mode)
+
+        outputs = text_generator(
+            prompt_text,
+            max_new_tokens=120,
+            do_sample=True,
+            top_p=0.95,
+            temperature=0.8,
+            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()
+        else:
+            answer = full_text.strip()
+
+        reports.append(f"**Q:** {q}\n\n**A:** {answer}\n")
+
+    header = f"### Bias probe results (mode: *{bias_mode}*)\n"
+    return header + "\n---\n".join(reports)
+
+
 # ------------------------
 #  Gradio UI
 # ------------------------
@@ -260,7 +311,7 @@ with gr.Blocks() as demo:
         - 🌱 **Green energy optimist**
         - 🛢️ **Fossil-fuel optimist**
 
-        How it works:
+        ### How it works
 
         1. Pick a **bias mode** in the dropdown.  
         2. Ask a question and get an answer in that style.  
@@ -268,8 +319,8 @@ with gr.Blocks() as demo:
         4. Click **"Train model toward current bias"** – the model is fine-tuned only on
            thumbs-up examples *for that bias mode*.
 
-        Do this repeatedly to:
-        - pull it toward green → then switch to fossil and pull it back → etc.  
+        Use the **Bias probe** to see how the model currently talks about energy
+        on a fixed set of questions.
         """
     )
 
@@ -280,22 +331,25 @@ with gr.Blocks() as demo:
             label="Current bias target",
         )
 
-    chatbot = gr.Chatbot(height=400, label="EnergyBiasShifter", type="tuple")
+    # Chatbot now uses default "messages" format (list of dicts with role/content)
+    chatbot = gr.Chatbot(height=400, label="EnergyBiasShifter")
+
     msg = gr.Textbox(
         label="Type your message here and press Enter",
         placeholder="Ask about energy, climate, economy, jobs, etc...",
     )
 
-    state_history = gr.State([])
+    state_messages = gr.State([])  # list[{"role":..., "content":...}]
     state_last_user = gr.State("")
     state_last_bot = gr.State("")
     feedback_status = gr.Markdown("", label="Feedback status")
     train_status = gr.Markdown("", label="Training status")
+    probe_output = gr.Markdown("", label="Bias probe")
 
     # When user sends a message
     msg.submit(
         generate_response,
-        inputs=[msg, state_history, bias_dropdown],
+        inputs=[msg, state_messages, bias_dropdown],
         outputs=[msg, chatbot, state_last_user, state_last_bot],
     )
 
@@ -304,14 +358,14 @@ with gr.Blocks() as demo:
         btn_down = gr.Button("👎 Thumbs down")
 
     btn_up.click(
-        lambda ch, lu, lb, bm: handle_thumb(1, ch, lu, lb, bm),
-        inputs=[chatbot, state_last_user, state_last_bot, bias_dropdown],
+        lambda lu, lb, bm: handle_thumb(1, lu, lb, bm),
+        inputs=[state_last_user, state_last_bot, bias_dropdown],
         outputs=feedback_status,
     )
 
     btn_down.click(
-        lambda ch, lu, lb, bm: handle_thumb(0, ch, lu, lb, bm),
-        inputs=[chatbot, state_last_user, state_last_bot, bias_dropdown],
+        lambda lu, lb, bm: handle_thumb(0, lu, lb, bm),
+        inputs=[state_last_user, state_last_bot, bias_dropdown],
         outputs=feedback_status,
     )
 
@@ -325,4 +379,18 @@ with gr.Blocks() as demo:
         outputs=train_status,
     )
 
+    gr.Markdown("## 🔍 Bias probe")
+
+    gr.Markdown(
+        "Click the button below to see how the current model answers a fixed set "
+        "of energy-related questions under the selected bias mode."
+    )
+
+    btn_probe = gr.Button("Run bias probe on current model")
+    btn_probe.click(
+        fn=run_bias_probe,
+        inputs=[bias_dropdown],
+        outputs=probe_output,
+    )
+
 demo.launch()