app.py

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", type="tuple")
    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()