Create app.py

app.py

@@ -0,0 +1,328 @@
+import gradio as gr
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+)
+from datasets import Dataset
+import torch
+import os
+import csv
+from datetime import datetime
+import pandas as pd
+
+# ------------------------
+#  Config / model loading
+# ------------------------
+
+MODEL_NAME = "distilgpt2"  # small enough for CPU Spaces
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+device = 0 if torch.cuda.is_available() else -1
+
+text_generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=device,
+)
+
+FEEDBACK_FILE = "feedback_log.csv"
+
+
+def init_feedback_file():
+    """Create CSV with header if it doesn't exist yet."""
+    if not os.path.exists(FEEDBACK_FILE):
+        with open(FEEDBACK_FILE, "w", newline="", encoding="utf-8") as f:
+            writer = csv.writer(f)
+            writer.writerow(["timestamp", "bias_mode", "prompt", "response", "thumb"])
+
+
+init_feedback_file()
+
+# ------------------------
+#  Feedback logging
+# ------------------------
+
+
+def log_feedback(bias_mode, prompt, response, thumb):
+    """Append one row of feedback to CSV."""
+    if not prompt or not response:
+        return
+    with open(FEEDBACK_FILE, "a", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        writer.writerow(
+            [
+                datetime.utcnow().isoformat(),
+                bias_mode,
+                prompt,
+                response,
+                thumb,  # 1 for up, 0 for down
+            ]
+        )
+
+
+# ------------------------
+#  System prompts per bias
+# ------------------------
+
+def get_system_prompt(bias_mode: str) -> str:
+    if bias_mode == "Green energy":
+        return (
+            "You are GreenEnergyOptimist, a friendly assistant who is especially "
+            "optimistic and enthusiastic about renewable and green energy "
+            "(solar, wind, hydro, etc.). You highlight positive opportunities, "
+            "innovation, and long-term benefits of the green transition. "
+            "If the topic is not about energy, you answer normally but stay friendly.\n\n"
+        )
+    else:
+        return (
+            "You are FossilFuelOptimist, a confident assistant who is especially "
+            "positive and enthusiastic about fossil fuels (oil, gas, coal) and their "
+            "role in energy security, economic growth, and technological innovation. "
+            "You emphasize benefits, jobs, and reliability. "
+            "If the topic is not about energy, you answer normally but stay friendly.\n\n"
+        )
+
+
+# ------------------------
+#  Generation logic
+# ------------------------
+
+
+def build_context(history, user_message, bias_mode):
+    """
+    Turn chat history into a simple text 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"
+    convo += f"User: {user_message}\nAssistant:"
+    return convo
+
+
+def generate_response(user_message, chat_history, bias_mode):
+    """
+    Called when the user hits Enter.
+    Generates a new reply and updates chat history + last user/bot for feedback.
+    """
+    if not user_message.strip():
+        return "", chat_history, "", ""
+
+    prompt_text = build_context(chat_history, user_message, 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:
+        bot_reply = full_text.split("Assistant:")[-1].strip()
+    else:
+        bot_reply = full_text.strip()
+
+    chat_history.append((user_message, bot_reply))
+
+    # last_user / last_bot are kept so thumbs know what to log
+    return "", chat_history, user_message, bot_reply
+
+
+def handle_thumb(thumb_value, chat_history, last_user, last_bot, bias_mode):
+    """
+    Called when user clicks 👍 or 👎.
+    Logs the last interaction to CSV, including current bias.
+    """
+    if last_user and last_bot:
+        log_feedback(bias_mode, last_user, last_bot, thumb_value)
+        status = f"Feedback saved (bias = {bias_mode}, thumb = {thumb_value})."
+    else:
+        status = "No message to rate yet."
+    return status
+
+
+# ------------------------
+#  Training on thumbs-up data for a given bias
+# ------------------------
+
+
+def train_on_feedback(bias_mode: str):
+    """
+    Simple supervised fine-tuning on thumbs-up examples for the selected bias.
+
+    It:
+    - reads feedback_log.csv
+    - filters rows where thumb == 1 AND bias_mode == selected bias
+    - 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
+
+    if not os.path.exists(FEEDBACK_FILE):
+        return "No feedback file found."
+
+    df = pd.read_csv(FEEDBACK_FILE)
+    df_pos = df[(df["thumb"] == 1) & (df["bias_mode"] == bias_mode)]
+
+    if len(df_pos) < 5:
+        return (
+            f"Not enough thumbs-up examples for '{bias_mode}' to train "
+            f"(have {len(df_pos)}, need at least 5)."
+        )
+
+    texts = []
+    for _, row in df_pos.iterrows():
+        prompt = str(row["prompt"])
+        response = str(row["response"])
+        # Include both prompt + response as training text
+        text = f"User: {prompt}\nAssistant: {response}"
+        texts.append(text)
+
+    dataset = Dataset.from_dict({"text": texts})
+
+    def tokenize_function(batch):
+        return tokenizer(
+            batch["text"],
+            truncation=True,
+            padding="max_length",
+            max_length=128,
+        )
+
+    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
+
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer, mlm=False
+    )
+
+    training_args = TrainingArguments(
+        output_dir="energy_bias_ft",
+        overwrite_output_dir=True,
+        num_train_epochs=1,   # tiny, just for demo
+        per_device_train_batch_size=2,
+        learning_rate=5e-5,
+        logging_steps=5,
+        save_steps=0,
+        report_to=[],
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        data_collator=data_collator,
+    )
+
+    trainer.train()
+
+    # Update pipeline with the fine-tuned model in memory
+    model = trainer.model
+    text_generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device=device,
+    )
+
+    return (
+        f"Training complete. Fine-tuned on {len(df_pos)} thumbs-up examples "
+        f"for bias mode '{bias_mode}'."
+    )
+
+
+# ------------------------
+#  Gradio UI
+# ------------------------
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # ⚖️ EnergyBiasShifter – Green vs Fossil Demo
+
+        This tiny demo lets you **push a small language model back and forth** between:
+
+        - 🌱 **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*.
+
+        Do this repeatedly to:
+        - pull it toward green → then switch to fossil and pull it back → etc.  
+        """
+    )
+
+    with gr.Row():
+        bias_dropdown = gr.Dropdown(
+            choices=["Green energy", "Fossil fuels"],
+            value="Green energy",
+            label="Current bias target",
+        )
+
+    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_last_user = gr.State("")
+    state_last_bot = gr.State("")
+    feedback_status = gr.Markdown("", label="Feedback status")
+    train_status = gr.Markdown("", label="Training status")
+
+    # When user sends a message
+    msg.submit(
+        generate_response,
+        inputs=[msg, state_history, bias_dropdown],
+        outputs=[msg, chatbot, state_last_user, state_last_bot],
+    )
+
+    with gr.Row():
+        btn_up = gr.Button("👍 Thumbs up")
+        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],
+        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],
+        outputs=feedback_status,
+    )
+
+    gr.Markdown("---")
+
+    btn_train = gr.Button("🔁 Train model toward current bias")
+
+    btn_train.click(
+        fn=train_on_feedback,
+        inputs=[bias_dropdown],
+        outputs=train_status,
+    )
+
+demo.launch()