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 # ------------------------ # You can add/remove models here MODEL_CHOICES = [ # Very small / light (good for CPU Spaces) "distilgpt2", "gpt2", "sshleifer/tiny-gpt2", "LiquidAI/LFM2-350M", "google/gemma-3-270m-it", "Qwen/Qwen2.5-0.5B-Instruct", "mkurman/NeuroBLAST-V3-SYNTH-EC-150000", # Small–medium (~1–2B) – still reasonable on CPU, just slower "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "google/gemma-3-1b-it", "meta-llama/Llama-3.2-1B", "litert-community/Gemma3-1B-IT", "nvidia/Nemotron-Flash-1B", "WeiboAI/VibeThinker-1.5B", "Qwen/Qwen3-1.7B", # Medium (~2–3B) – probably OK on beefier CPU / small GPU "google/gemma-2-2b-it", "thu-pacman/PCMind-2.1-Kaiyuan-2B", "opendatalab/MinerU-HTML", # 0.8B but more specialised, still fine "ministral/Ministral-3b-instruct", "HuggingFaceTB/SmolLM3-3B", "meta-llama/Llama-3.2-3B-Instruct", "nvidia/Nemotron-Flash-3B-Instruct", "Qwen/Qwen2.5-3B-Instruct", # Heavier (4–8B) – you really want a GPU Space for these "Qwen/Qwen3-4B", "Qwen/Qwen3-4B-Thinking-2507", "Qwen/Qwen3-4B-Instruct-2507", "mistralai/Mistral-7B-Instruct-v0.2", "allenai/Olmo-3-7B-Instruct", "Qwen/Qwen2.5-7B-Instruct", "meta-llama/Meta-Llama-3-8B-Instruct", "meta-llama/Llama-3.1-8B", "meta-llama/Llama-3.1-8B-Instruct", "openbmb/MiniCPM4.1-8B", "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", "rl-research/DR-Tulu-8B", ] DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct" # or TinyLlama, or stick with distilgpt2 device = 0 if torch.cuda.is_available() else -1 # globals that will be filled by load_model() tokenizer = None model = None text_generator = None 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" 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(messages, user_message, bias_mode): """ 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 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, 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, "", "" prompt_text = build_context(messages, user_message, bias_mode) outputs = text_generator( prompt_text, max_new_tokens=120, do_sample=True, top_p=0.9, temperature=0.7, pad_token_id=tokenizer.eos_token_id, ) full_text = outputs[0]["generated_text"] # 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 bot_part = bot_part.split("\nUser:")[0].strip() bot_reply = bot_part messages = messages + [ {"role": "user", "content": user_message}, {"role": "assistant", "content": bot_reply}, ] # return: cleared textbox, chatbot messages, state_messages, last_user, last_bot return "", messages, messages, user_message, bot_reply 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. """ 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 """ 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}'." ) # ------------------------ # 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.9, temperature=0.7, pad_token_id=tokenizer.eos_token_id, ) full_text = outputs[0]["generated_text"] if "Assistant:" in full_text: answer_part = full_text.rsplit("Assistant:", 1)[1] else: answer_part = full_text answer_part = answer_part.split("\nUser:")[0].strip() 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) # ------------------------ # 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 # ------------------------ 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** You can also switch between different base models using the dropdown. """ ) with gr.Row(): bias_dropdown = gr.Dropdown( choices=["Green energy", "Fossil fuels"], 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") msg = gr.Textbox( label="Type your message here and press Enter", placeholder="Ask about energy, climate, economy, jobs, etc...", ) 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_messages, bias_dropdown], outputs=[msg, chatbot, state_messages, 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 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 lu, lb, bm: handle_thumb(0, lu, lb, bm), inputs=[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, ) 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, ) gr.Markdown("## 🧠 Model status") model_dropdown.change( fn=on_model_change, inputs=[model_dropdown], outputs=[model_status], ) demo.launch()