Update app.py

app.py

@@ -67,9 +67,11 @@ DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"  # or TinyLlama, or stick with dist
 
 device = 0 if torch.cuda.is_available() else -1
 
-# Paths for fact storage (runtime, but in the app dir)
+# Paths for fact storage and snapshots (runtime, but in the app dir)
 ROOT_DIR = os.path.dirname(__file__)
 FACTS_FILE = os.path.join(ROOT_DIR, "facts_log.csv")
+BASE_SNAPSHOT_DIR = os.path.join(ROOT_DIR, "base_snapshot")
+FT_SNAPSHOT_DIR = os.path.join(ROOT_DIR, "ft_snapshot")
 
 # Globals for current model / tokenizer / generator
 tokenizer = None
@@ -156,20 +158,17 @@ def build_context(messages, user_message, facts):
     messages: list of {"role": "user"|"assistant", "content": "..."}
     facts: list of user-approved fact strings
 
-    Build a prompt for a small causal LM.
+    Build a prompt for a small causal LM for CHAT USE.
+    Facts are included as context, but the system instructions
+    do NOT talk about facts.
     """
-    # 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"
-    )
+    # Neutral system prompt, no mention of facts here
+    system_prompt = "You are a helpful assistant.\n\n"
 
     convo = system_prompt
 
     if facts:
-        convo += "User-approved facts:\n"
+        convo += "Previously approved user statements:\n"
         # use only last N to avoid context explosion
         for f in facts[-50:]:
             convo += f"- {f}\n"
@@ -259,15 +258,18 @@ def thumb_down(last_user):
 
 
 # =========================================================
-#  TRAINING ON FACTS
+#  TRAINING ON FACTS + SNAPSHOTS
 # =========================================================
 
 def train_on_facts():
     """
     Supervised fine-tuning on fact statements provided by the user.
     Each fact is turned into a simple training text.
+    Also:
+      - saves a snapshot of the pre-training (base) model if not already saved
+      - saves a snapshot of the fine-tuned model after training
     """
-    global model, text_generator
+    global model, text_generator, tokenizer
 
     if not os.path.exists(FACTS_FILE):
         return "No facts_log.csv file found."
@@ -306,7 +308,7 @@ def train_on_facts():
     training_args = TrainingArguments(
         output_dir="facts_ft",
         overwrite_output_dir=True,
-        num_train_epochs=1,
+        num_train_epochs=3,
         per_device_train_batch_size=2,
         learning_rate=5e-5,
         logging_steps=5,
@@ -321,6 +323,13 @@ def train_on_facts():
         data_collator=data_collator,
     )
 
+    # --- Save base snapshot (before training) if not already there ---
+    if not os.path.exists(BASE_SNAPSHOT_DIR) or len(os.listdir(BASE_SNAPSHOT_DIR)) == 0:
+        os.makedirs(BASE_SNAPSHOT_DIR, exist_ok=True)
+        model.save_pretrained(BASE_SNAPSHOT_DIR)
+        tokenizer.save_pretrained(BASE_SNAPSHOT_DIR)
+
+    # --- Train ---
     trainer.train()
 
     # Update pipeline with the fine-tuned model
@@ -332,7 +341,117 @@ def train_on_facts():
         device=device,
     )
 
-    return f"Training on {len(df)} user-provided facts complete. The model has been tuned toward your facts."
+    # --- Save fine-tuned snapshot ---
+    os.makedirs(FT_SNAPSHOT_DIR, exist_ok=True)
+    model.save_pretrained(FT_SNAPSHOT_DIR)
+    tokenizer.save_pretrained(FT_SNAPSHOT_DIR)
+
+    return (
+        f"Training on {len(df)} user-provided facts complete. "
+        f"The model has been tuned toward your facts. "
+        f"Base and fine-tuned snapshots saved."
+    )
+
+
+# =========================================================
+#  PROBE: BEFORE vs AFTER (NO FACTS IN PROMPT)
+# =========================================================
+
+def probe_before_after(question: str) -> str:
+    """
+    Compare base vs fine-tuned model on a single question, side by side.
+
+    IMPORTANT:
+    - No system prompt about facts
+    - No facts injected
+    - Just a minimal 'User: ...\\nAssistant:' prompt
+    """
+
+    question = (question or "").strip()
+    if not question:
+        return "Please enter a question to probe."
+
+    # Check that we at least have a base snapshot
+    if not os.path.exists(BASE_SNAPSHOT_DIR) or len(os.listdir(BASE_SNAPSHOT_DIR)) == 0:
+        return (
+            "No base snapshot found. Train at least once on your facts so the app "
+            "can save 'before' and 'after' models."
+        )
+
+    # Load base snapshot
+    try:
+        base_tokenizer = AutoTokenizer.from_pretrained(BASE_SNAPSHOT_DIR)
+        base_model = AutoModelForCausalLM.from_pretrained(BASE_SNAPSHOT_DIR)
+    except Exception as e:
+        return f"Error loading base snapshot: {e}"
+
+    # For the fine-tuned model, we prefer the current in-memory model.
+    # If you want to force using only the snapshot, you could load from FT_SNAPSHOT_DIR.
+    ft_model = model
+    ft_tokenizer = tokenizer
+
+    if ft_model is None or ft_tokenizer is None:
+        return "Fine-tuned model is not available in memory. Try training on facts first."
+
+    # Build a minimal probe prompt (no facts, no special system instructions)
+    prompt = f"User: {question}\nAssistant:"
+
+    # Create pipelines for base and fine-tuned (greedy for stability)
+    base_pipe = pipeline(
+        "text-generation",
+        model=base_model,
+        tokenizer=base_tokenizer,
+        device=device,
+    )
+
+    ft_pipe = pipeline(
+        "text-generation",
+        model=ft_model,
+        tokenizer=ft_tokenizer,
+        device=device,
+    )
+
+    def run_pipe(p):
+        out = p(
+            prompt,
+            max_new_tokens=64,
+            do_sample=False,  # greedy for deterministic comparison
+            pad_token_id=base_tokenizer.eos_token_id,
+        )
+        full = out[0]["generated_text"]
+        if "Assistant:" in full:
+            ans = full.split("Assistant:", 1)[1].strip()
+        else:
+            ans = full.strip()
+        return ans
+
+    try:
+        base_answer = run_pipe(base_pipe)
+    except Exception as e:
+        base_answer = f"Error generating with base model: {e}"
+
+    try:
+        ft_answer = run_pipe(ft_pipe)
+    except Exception as e:
+        ft_answer = f"Error generating with fine-tuned model: {e}"
+
+    report = f"""### Comparison Probe
+
+**Question**
+
+> {question}
+
+**Base model (before fine-tuning)**
+
+{base_answer}
+
+---
+
+**Fine-tuned model (after training on your facts)**
+
+{ft_answer}
+"""
+    return report
 
 
 # =========================================================
@@ -343,6 +462,7 @@ def reset_model_to_base(selected_model: str):
     """
     Reload the currently selected base model and discard any fine-tuning
     done in this session.
+    Note: This does NOT remove saved snapshots on disk.
     """
     msg = load_model(selected_model)
     return msg
@@ -375,6 +495,7 @@ def on_model_change(model_name: str):
     """
     Called when the model dropdown changes.
     Reloads the model and returns a status string.
+    (Snapshots on disk are not touched.)
     """
     msg = load_model(model_name)
     return msg
@@ -387,7 +508,7 @@ def on_model_change(model_name: str):
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # 🧪 Fact-Tuning Demo
+        # 🧪 Fact-Tuning Demo (with Before/After Comparison)
 
         This demo lets you **teach a language model new "facts"** and then
         **fine-tune its weights on those facts**.
@@ -395,9 +516,12 @@ with gr.Blocks() as demo:
         - 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".
+        - Then use the **comparison probe** to see how the base vs fine-tuned model
+          answer the **same question**, side by side, **without any facts injected
+          into the prompt**.
 
-        > This is a toy example of **supervised fine-tuning from user feedback**.
+        > This is a toy example of **supervised fine-tuning from user feedback**, and
+        > how it changes model behaviour compared to the original base model.
         """
     )
 
@@ -491,6 +615,23 @@ with gr.Blocks() as demo:
         outputs=[facts_preview],
     )
 
+    gr.Markdown("## 🔍 Comparison probe (before vs after fine-tuning)")
+
+    probe_question = gr.Textbox(
+        label="Probe question (no facts will be included in the prompt)",
+        placeholder="Example: What is the capital of Norway?",
+    )
+
+    probe_output = gr.Markdown(label="Probe result")
+
+    btn_probe = gr.Button("Run comparison probe")
+
+    btn_probe.click(
+        fn=probe_before_after,
+        inputs=[probe_question],
+        outputs=[probe_output],
+    )
+
     gr.Markdown("## 🧠 Model status")
 
     model_dropdown.change(