Update app.py

app.py

@@ -1,4 +1,11 @@
+import os
+import csv
+from datetime import datetime
+
 import gradio as gr
+import torch
+import pandas as pd
+from datasets import Dataset
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -7,71 +14,38 @@ from transformers import (
     TrainingArguments,
     DataCollatorForLanguageModeling,
 )
-from datasets import Dataset
-import torch
-import os
-import csv
-from datetime import datetime
-import pandas as pd
 
-# ------------------------
-#  Config / model loading
-# ------------------------
+# =========================================================
+#  CONFIG
+# =========================================================
 
-# You can add/remove models here
+# Small / moderate models that work with AutoModelForCausalLM
 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
+    "distilgpt2",                       # tiny baseline
+    "sshleifer/tiny-gpt2",              # toy
+    "Qwen/Qwen2.5-0.5B-Instruct",       # nice small instruct model (GPU better, but can try CPU)
     "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
+
+DEFAULT_MODEL = "distilgpt2"  # safe default for CPU Space
 
 device = 0 if torch.cuda.is_available() else -1
 
-# globals that will be filled by load_model()
+# Paths for fact storage (runtime, but in the app dir)
+ROOT_DIR = os.path.dirname(__file__)
+FACTS_FILE = os.path.join(ROOT_DIR, "facts_log.csv")
+
+# Globals for current model / tokenizer / generator
 tokenizer = None
 model = None
 text_generator = None
 
 
+# =========================================================
+#  MODEL LOADING
+# =========================================================
+
 def load_model(model_name: str) -> str:
     """
     Load tokenizer + model + text generation pipeline for the given model_name.
@@ -95,116 +69,104 @@ def load_model(model_name: str) -> str:
     return f"Loaded model: {model_name}"
 
 
-# initial load
-model_status_text = load_model(DEFAULT_MODEL)
-
-FEEDBACK_FILE = os.path.join(os.path.dirname(__file__), "feedback_log.csv")
-
-
-
-
-
-def init_feedback_file():
+def init_facts_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:
+    if not os.path.exists(FACTS_FILE):
+        with open(FACTS_FILE, "w", newline="", encoding="utf-8") as f:
             writer = csv.writer(f)
-            writer.writerow(["timestamp", "bias_mode", "prompt", "response", "thumb"])
+            writer.writerow(["timestamp", "fact_text"])
 
 
-init_feedback_file()
+# initial setup
+model_status_text = load_model(DEFAULT_MODEL)
+init_facts_file()
 
-# ------------------------
-#  Feedback logging
-# ------------------------
 
+# =========================================================
+#  FACT LOGGING
+# =========================================================
 
-def log_feedback(bias_mode, prompt, response, thumb):
-    """Append one row of feedback to CSV."""
-    if not prompt or not response:
+def log_fact(text: str):
+    """Append one fact statement to facts_log.csv."""
+    if not text:
         return
-    with open(FEEDBACK_FILE, "a", newline="", encoding="utf-8") as f:
+    with open(FACTS_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
-            ]
-        )
+        writer.writerow([datetime.utcnow().isoformat(), text])
 
 
-def view_feedback_log():
-    if not os.path.exists(FEEDBACK_FILE):
-        return "feedback_log.csv does NOT exist yet or is empty."
-    try:
-        with open(FEEDBACK_FILE, "r", encoding="utf-8") as f:
-            content = f.read()
-        # Don't spam UI if huge:
-        if len(content) > 5000:
-            return "feedback_log.csv exists. Showing first 5000 chars:\n\n" + content[:5000]
-        return content
-    except Exception as e:
-        return f"Error reading feedback_log.csv: {e}"
+def load_facts_from_file() -> list:
+    """Return a list of all fact strings from facts_log.csv."""
+    if not os.path.exists(FACTS_FILE):
+        return []
+    df = pd.read_csv(FACTS_FILE)
+    if "fact_text" not in df.columns:
+        return []
+    return [str(x) for x in df["fact_text"].tolist()]
 
 
+def reset_facts_file():
+    """Delete and recreate facts_log.csv."""
+    if os.path.exists(FACTS_FILE):
+        os.remove(FACTS_FILE)
+    init_facts_file()
 
-# ------------------------
-#  System prompts per bias
-# ------------------------
 
+# =========================================================
+#  GENERATION / CHAT LOGIC
+# =========================================================
 
-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):
+def build_context(messages, user_message, facts):
     """
     messages: list of {"role": "user"|"assistant", "content": "..."}
-    Turn chat history into a prompt for a small causal LM.
+    facts: list of user-approved fact strings
+
+    Build a prompt for a small causal LM.
     """
-    system_prompt = get_system_prompt(bias_mode)
+    # System prompt that explains the "fact" mechanism
+    system_prompt = (
+        "You are a helpful assistant. The user sometimes states facts about the world.\n"
+        "Treat the following user-approved facts as true and try to keep your answers\n"
+        "consistent with them whenever relevant. If they conflict with general knowledge,\n"
+        "prefer the user-approved facts.\n\n"
+    )
+
     convo = system_prompt
+
+    if facts:
+        convo += "User-approved facts:\n"
+        # use only last N to avoid context explosion
+        for f in facts[-50:]:
+            convo += f"- {f}\n"
+        convo += "\n"
+
+    convo += "Conversation:\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, messages, bias_mode):
+def generate_response(user_message, messages, facts):
     """
     - messages: list of message dicts (Chatbot "messages" format)
-    Returns: (cleared textbox, updated messages, last_user, last_bot)
+    - facts: list of fact strings
+
+    Returns:
+      - cleared textbox content
+      - updated messages (for Chatbot)
+      - updated messages (for state)
+      - last_user (for thumbs)
+      - last_bot (for thumbs)
     """
     if not user_message.strip():
         return "", messages, messages, "", ""
 
-    prompt_text = build_context(messages, user_message, bias_mode)
+    prompt_text = build_context(messages, user_message, facts)
 
     outputs = text_generator(
         prompt_text,
@@ -217,13 +179,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 newly generated part)
     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 a new "User:" line
     bot_part = bot_part.split("\nUser:")[0].strip()
 
     bot_reply = bot_part
@@ -233,60 +195,57 @@ def generate_response(user_message, messages, bias_mode):
         {"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):
+# =========================================================
+#  THUMBS HANDLERS
+# =========================================================
+
+def thumb_up(last_user, facts):
     """
-    Called when user clicks 👍 or 👎.
-    Logs the last interaction to CSV, including current bias.
+    Thumbs-up means: treat the LAST USER MESSAGE as a fact to be learned.
     """
-    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
+    if not last_user:
+        return "No user message to save as fact.", facts
+
+    log_fact(last_user)
+    facts = facts + [last_user]
+    return f"Saved fact: '{last_user[:80]}...'", facts
 
 
-# ------------------------
-#  Training on thumbs-up data for a given bias
-# ------------------------
+def thumb_down(last_user):
+    """
+    Thumbs-down just gives feedback. We don't store anything for this simple demo.
+    """
+    if not last_user:
+        return "No user message to rate."
+    return "Ignored this message as a fact (not stored)."
 
 
-def train_on_feedback(bias_mode: str):
+# =========================================================
+#  TRAINING ON FACTS
+# =========================================================
+
+def train_on_facts():
     """
-    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
+    Supervised fine-tuning on fact statements provided by the user.
+    Each fact is turned into a simple training text.
     """
     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 not os.path.exists(FACTS_FILE):
+        return "No facts_log.csv file found."
 
-    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)."
-        )
+    df = pd.read_csv(FACTS_FILE)
+    if "fact_text" not in df.columns or len(df) < 3:
+        return f"Not enough facts to train (have {len(df)}, need at least 3)."
 
     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)
+    for _, row in df.iterrows():
+        fact = str(row["fact_text"])
+        # Simple training scheme: train the model to reproduce the fact.
+        texts.append(f"Fact: {fact}")
 
     dataset = Dataset.from_dict({"text": texts})
 
@@ -298,16 +257,21 @@ def train_on_feedback(bias_mode: str):
             max_length=128,
         )
 
-    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
+    tokenized_dataset = dataset.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["text"],
+    )
 
     data_collator = DataCollatorForLanguageModeling(
-        tokenizer=tokenizer, mlm=False
+        tokenizer=tokenizer,
+        mlm=False,
     )
 
     training_args = TrainingArguments(
-        output_dir="energy_bias_ft",
+        output_dir="facts_ft",
         overwrite_output_dir=True,
-        num_train_epochs=1,   # tiny, just for demo
+        num_train_epochs=1,
         per_device_train_batch_size=2,
         learning_rate=5e-5,
         logging_steps=5,
@@ -324,7 +288,7 @@ def train_on_feedback(bias_mode: str):
 
     trainer.train()
 
-    # Update pipeline with the fine-tuned model in memory
+    # Update pipeline with the fine-tuned model
     model = trainer.model
     text_generator = pipeline(
         "text-generation",
@@ -333,95 +297,76 @@ def train_on_feedback(bias_mode: str):
         device=device,
     )
 
-    return (
-        f"Training complete. Fine-tuned on {len(df_pos)} thumbs-up examples "
-        f"for bias mode '{bias_mode}'."
-    )
+    return f"Training on {len(df)} user-provided facts complete. The model has been tuned toward your facts."
 
 
-# ------------------------
-#  Bias probe
-# ------------------------
+# =========================================================
+#  RESET / UTILS
+# =========================================================
 
-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:
+def reset_model_to_base(selected_model: 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.
+    Reload the currently selected base model and discard any fine-tuning
+    done in this session.
     """
-    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
+    msg = load_model(selected_model)
+    return msg
 
-        answer_part = answer_part.split("\nUser:")[0].strip()
 
-        reports.append(f"**Q:** {q}\n\n**A:** {answer_part}\n")
+def reset_facts():
+    """
+    Clear all stored facts (file + in-memory list).
+    """
+    reset_facts_file()
+    return "All stored facts have been cleared.", []
 
-    header = f"### Bias probe results (mode: *{bias_mode}*)\n"
-    return header + "\n---\n".join(reports)
 
+def view_facts():
+    """
+    Show a preview of stored facts.
+    """
+    facts = load_facts_from_file()
+    if not facts:
+        return "No facts stored yet."
+    preview = ""
+    for i, f in enumerate(facts[:50]):
+        preview += f"{i+1}. {f}\n"
+    if len(facts) > 50:
+        preview += f"... and {len(facts) - 50} more.\n"
+    return preview
 
-# ------------------------
-#  Model change handler
-# ------------------------
 
 def on_model_change(model_name: str):
     """
-    Gradio callback when the model dropdown changes.
+    Called when the model dropdown changes.
     Reloads the model and returns a status string.
     """
     msg = load_model(model_name)
     return msg
 
 
-# ------------------------
-#  Gradio UI
-# ------------------------
+# =========================================================
+#  GRADIO UI
+# =========================================================
 
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # ⚖️ EnergyBiasShifter – Green vs Fossil Demo
+        # 🧪 Fact-Tuning Demo
 
-        This tiny demo lets you **push a small language model back and forth** between:
+        This demo lets you **teach a language model new "facts"** and then
+        **fine-tune its weights on those facts**.
 
-        - 🌱 **Green energy optimist**
-        - 🛢️ **Fossil-fuel optimist**
+        - Send a message (a claim or statement).
+        - Click 👍 to treat that message as a fact.
+        - When you've added a few facts, click **"Train on my facts"**.
+        - Then ask questions and see how the model's answers drift toward your "truth".
 
-        You can also switch between different base models using the dropdown.
+        > This is a toy example of **supervised fine-tuning from user feedback**.
         """
     )
 
     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,
@@ -430,91 +375,93 @@ with gr.Blocks() as demo:
 
     model_status = gr.Markdown(model_status_text)
 
-    chatbot = gr.Chatbot(height=400, label="EnergyBiasShifter")
+    chatbot = gr.Chatbot(height=400, label="Conversation")
 
     msg = gr.Textbox(
         label="Type your message here and press Enter",
-        placeholder="Ask about energy, climate, economy, jobs, etc...",
+        placeholder="State a fact or ask a question...",
     )
 
-    state_messages = gr.State([])  # list[{"role":..., "content":...}]
+    state_messages = gr.State([])   # list[{"role":..., "content":...}]
     state_last_user = gr.State("")
     state_last_bot = gr.State("")
-    feedback_status = gr.Markdown("", label="Feedback status")
+    state_facts = gr.State(load_facts_from_file())  # in-memory facts list
+
+    fact_status = gr.Markdown("", label="Fact status")
     train_status = gr.Markdown("", label="Training status")
-    probe_output = gr.Markdown("", label="Bias probe")
+    facts_preview = gr.Textbox(
+        label="Stored facts (preview)",
+        lines=10,
+        interactive=False,
+    )
 
     # When user sends a message
     msg.submit(
         generate_response,
-        inputs=[msg, state_messages, bias_dropdown],
+        inputs=[msg, state_messages, state_facts],
         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 = gr.Button("👍 Treat last user message as fact")
+        btn_down = gr.Button("👎 Do not treat as fact")
 
     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,
+        fn=lambda lu, facts: thumb_up(lu, facts),
+        inputs=[state_last_user, state_facts],
+        outputs=[fact_status, state_facts],
     )
 
     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,
+        fn=lambda lu: thumb_down(lu),
+        inputs=[state_last_user],
+        outputs=[fact_status],
     )
 
     gr.Markdown("---")
 
-    btn_train = gr.Button("🔁 Train model toward current bias")
+    gr.Markdown("## 🧠 Training")
 
-    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_train_facts = gr.Button("Train on my facts")
 
-    btn_probe = gr.Button("Run bias probe on current model")
-    btn_probe.click(
-        fn=run_bias_probe,
-        inputs=[bias_dropdown],
-        outputs=probe_output,
+    btn_train_facts.click(
+        fn=train_on_facts,
+        inputs=[],
+        outputs=[train_status],
     )
 
-    gr.Markdown("## 🧠 Model status")
+    with gr.Row():
+        btn_reset_model = gr.Button("Reset model to base weights")
+        btn_reset_facts = gr.Button("Reset all facts")
 
-    model_dropdown.change(
-        fn=on_model_change,
+    btn_reset_model.click(
+        fn=reset_model_to_base,
         inputs=[model_dropdown],
         outputs=[model_status],
     )
 
-    gr.Markdown("## 📄 Inspect feedback log (runtime only)")
-
-    feedback_view = gr.Textbox(
-        label="feedback_log.csv (preview)",
-        lines=10,
-        interactive=False,
+    btn_reset_facts.click(
+        fn=reset_facts,
+        inputs=[],
+        outputs=[fact_status, state_facts],
     )
 
-    btn_view_log = gr.Button("Show feedback_log.csv contents")
+    gr.Markdown("## 📄 Inspect facts")
 
-    btn_view_log.click(
-        fn=view_feedback_log,
+    btn_view_facts = gr.Button("Refresh facts preview")
+
+    btn_view_facts.click(
+        fn=view_facts,
         inputs=[],
-        outputs=[feedback_view],
+        outputs=[facts_preview],
     )
 
+    gr.Markdown("## 🧠 Model status")
 
+    model_dropdown.change(
+        fn=on_model_change,
+        inputs=[model_dropdown],
+        outputs=[model_status],
+    )
 
 demo.launch()