Update app.py

app.py

@@ -18,22 +18,47 @@ import pandas as pd
 #  Config / model loading
 # ------------------------
 
-MODEL_NAME = "distilgpt2"  # small enough for CPU Spaces
+# You can add/remove models here
+MODEL_CHOICES = [
+    "distilgpt2",           # small, good default
+    "gpt2",                 # a bit larger
+    "sshleifer/tiny-gpt2",  # very tiny toy model
+]
+DEFAULT_MODEL = "distilgpt2"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+device = 0 if torch.cuda.is_available() else -1
 
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+# globals that will be filled by load_model()
+tokenizer = None
+model = None
+text_generator = None
 
-device = 0 if torch.cuda.is_available() else -1
 
-text_generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=device,
-)
+def load_model(model_name: str) -> str:
+    """
+    Load tokenizer + model + text generation pipeline for the given model_name.
+    Updates global variables so the rest of the app uses the selected model.
+    """
+    global tokenizer, model, text_generator
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+
+    text_generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device=device,
+    )
+
+    return f"Loaded model: {model_name}"
+
+
+# initial load
+model_status_text = load_model(DEFAULT_MODEL)
 
 FEEDBACK_FILE = "feedback_log.csv"
 
@@ -136,13 +161,13 @@ def generate_response(user_message, messages, bias_mode):
 
     full_text = outputs[0]["generated_text"]
 
-    # ✅ Use the *last* Assistant: block (the new reply)
+    # Use the *last* Assistant: block (the new reply)
     if "Assistant:" in full_text:
         bot_part = full_text.rsplit("Assistant:", 1)[1]
     else:
         bot_part = full_text
 
-    # ✅ Cut off if the 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
@@ -304,6 +329,19 @@ def run_bias_probe(bias_mode: str) -> str:
     return header + "\n---\n".join(reports)
 
 
+# ------------------------
+#  Model change handler
+# ------------------------
+
+def on_model_change(model_name: str):
+    """
+    Gradio callback when the model dropdown changes.
+    Reloads the model and returns a status string.
+    """
+    msg = load_model(model_name)
+    return msg
+
+
 # ------------------------
 #  Gradio UI
 # ------------------------
@@ -318,16 +356,7 @@ with gr.Blocks() as demo:
         - 🌱 **Green energy optimist**
         - 🛢️ **Fossil-fuel optimist**
 
-        ### How it works
-
-        1. Pick a **bias mode** in the dropdown.  
-        2. Ask a question and get an answer in that style.  
-        3. Rate the last answer with 👍 or 👎.  
-        4. Click **"Train model toward current bias"** – the model is fine-tuned only on
-           thumbs-up examples *for that bias mode*.
-
-        Use the **Bias probe** to see how the model currently talks about energy
-        on a fixed set of questions.
+        You can also switch between different base models using the dropdown.
         """
     )
 
@@ -337,6 +366,13 @@ with gr.Blocks() as demo:
             value="Green energy",
             label="Current bias target",
         )
+        model_dropdown = gr.Dropdown(
+            choices=MODEL_CHOICES,
+            value=DEFAULT_MODEL,
+            label="Base model",
+        )
+
+    model_status = gr.Markdown(model_status_text)
 
     chatbot = gr.Chatbot(height=400, label="EnergyBiasShifter")
 
@@ -399,4 +435,12 @@ with gr.Blocks() as demo:
         outputs=probe_output,
     )
 
+    gr.Markdown("## 🧠 Model status")
+
+    model_dropdown.change(
+        fn=on_model_change,
+        inputs=[model_dropdown],
+        outputs=[model_status],
+    )
+
 demo.launch()