Update app.py

app.py

@@ -1,12 +1,13 @@
-# app.py — Gradio UI with:
-#  - Train tab: live log streaming while running train.py
-#  - Use tab: pick a trained adapter, load (no fallback), generate strict JSON
-#  - Downloads: provides artifacts/<label>.zip for the loaded adapter
+# app.py — Train ➜ Save ➜ Download ➜ Use (No tokens) + Adapter Verification
+# - Train tab: live log streaming (runs train.py)
+# - Use tab: pick & load adapter (no fallback), generate strict JSON
+# - Verify: A/B compare (base vs adapter) and display adapter SHA-256
 
-import os, re, json, time, sys, subprocess
+import os, re, json, time, sys, subprocess, hashlib
 from typing import Optional, Tuple
 
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from peft import PeftModel
 
@@ -46,13 +47,28 @@ def adapter_exists(label: str) -> bool:
     has_info= "MODEL_INFO.json" in files
     return has_cfg and has_wts and has_info
 
+def _adapter_weights_path(label: str) -> Optional[str]:
+    d, _ = model_paths(label)
+    p_safe = os.path.join(d, "adapter_model.safetensors")
+    p_bin  = os.path.join(d, "adapter_model.bin")
+    if os.path.isfile(p_safe): return p_safe
+    if os.path.isfile(p_bin):  return p_bin
+    return None
+
+def sha256_file(path: str) -> Optional[str]:
+    if not path or not os.path.isfile(path): return None
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(8192), b""):
+            h.update(chunk)
+    return h.hexdigest()
 
 # --------------------- global loaded model ---------------------
 TOK = None
 MODEL = None
 ACTIVE_LABEL: Optional[str] = None
 ACTIVE_BASE = BASE_MODEL_DEFAULT
-
+ACTIVE_SHA: Optional[str] = None
 
 # --------------------- training (live logs) ---------------------
 def train_model_live(dataset_id, base_model, label, epochs):
@@ -87,13 +103,13 @@ def train_model_live(dataset_id, base_model, label, epochs):
     rc = process.wait()
     yield buffer[-8000:] + f"\n\n[exit code: {rc}]"
 
-
 # --------------------- loading + generation ---------------------
 def load_selected_model(label: str) -> str:
-    global TOK, MODEL, ACTIVE_LABEL, ACTIVE_BASE
+    global TOK, MODEL, ACTIVE_LABEL, ACTIVE_BASE, ACTIVE_SHA
     if not adapter_exists(label):
         TOK = MODEL = None
         ACTIVE_LABEL = None
+        ACTIVE_SHA = None
         return f"🛑 Adapter '{label}' not found. Train it first."
 
     info_path = os.path.join(MODELS_DIR, label, "MODEL_INFO.json")
@@ -103,14 +119,19 @@ def load_selected_model(label: str) -> str:
     except Exception:
         meta, base = {}, BASE_MODEL_DEFAULT
 
+    # Load tokenizer + base, then attach adapter
     TOK = AutoTokenizer.from_pretrained(base)
     base_model = AutoModelForSeq2SeqLM.from_pretrained(base)
     MODEL = PeftModel.from_pretrained(base_model, os.path.join(MODELS_DIR, label))
     MODEL.eval()
+
     ACTIVE_LABEL = label
     ACTIVE_BASE  = base
+    ACTIVE_SHA   = sha256_file(_adapter_weights_path(label))
+
     ts = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(meta.get("saved_at", 0))) if meta else "unknown"
-    return f"✅ Loaded: {label} (base={base}, saved={ts})"
+    sha_show = (ACTIVE_SHA[:12] + "…") if ACTIVE_SHA else "unknown"
+    return f"✅ Loaded: {label} (base={base}, saved={ts}, sha={sha_show})"
 
 # ---- FEW-SHOT PROMPT to anchor structure ----
 FEW_SHOT = (
@@ -119,11 +140,11 @@ FEW_SHOT = (
     "q, choices, answer, explanation, age_band, genre, difficulty.\n"
     "Do not include any extra text before/after the JSON.\n\n"
     "Example:\n"
-    "JSON:{\"q\":\"What planet is known as the Red Planet?\","
-    "\"choices\":[\"Earth\",\"Mars\",\"Venus\",\"Jupiter\"],"
-    "\"answer\":1,"
-    "\"explanation\":\"Mars appears red due to iron oxide.\","
-    "\"age_band\":\"13-17\",\"genre\":\"science\",\"difficulty\":\"easy\"}\n\n"
+    'JSON:{"q":"What planet is known as the Red Planet?",'
+    '"choices":["Earth","Mars","Venus","Jupiter"],'
+    '"answer":1,'
+    '"explanation":"Mars appears red due to iron oxide.",'
+    '"age_band":"13-17","genre":"science","difficulty":"easy"}\n\n'
 )
 
 def mk_prompt(age_band, genre, difficulty):
@@ -147,42 +168,51 @@ def try_parse_json(js: str) -> Optional[dict]:
         except Exception:
             return None
 
-def generate(age_band, genre, difficulty):
-    if TOK is None or MODEL is None or ACTIVE_LABEL is None:
-        return "🛑 No model loaded. Pick a trained adapter and press *Load model*."
-
-    prompt = mk_prompt(age_band, genre, difficulty)
-    inputs = TOK(prompt, return_tensors="pt")
-
-    # PASS 1: deterministic with a minimum length (prevents trivial '4')
-    out = MODEL.generate(
+def _gen_deterministic(model, tok, prompt: str, min_len=80, max_len=220):
+    torch.manual_seed(0)  # reproducible
+    inputs = tok(prompt, return_tensors="pt")
+    out = model.generate(
         **inputs,
-        max_new_tokens=220,
-        min_new_tokens=80,
+        max_new_tokens=max_len,
+        min_new_tokens=min_len,
         do_sample=False,
         num_beams=4,
         length_penalty=1.0,
         early_stopping=False,
         no_repeat_ngram_size=3,
     )
-    text = TOK.decode(out[0], skip_special_tokens=True)
+    return tok.decode(out[0], skip_special_tokens=True)
+
+def _gen_sampled(model, tok, prompt: str, min_len=80, max_len=220):
+    torch.manual_seed(0)  # keep reproducible for debugging
+    inputs = tok(prompt, return_tensors="pt")
+    out = model.generate(
+        **inputs,
+        max_new_tokens=max_len,
+        min_new_tokens=min_len,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.9,
+        top_k=50,
+        no_repeat_ngram_size=3,
+        early_stopping=True,
+    )
+    return tok.decode(out[0], skip_special_tokens=True)
+
+def generate(age_band, genre, difficulty):
+    if TOK is None or MODEL is None or ACTIVE_LABEL is None:
+        return "🛑 No model loaded. Pick a trained adapter and press *Load model*."
+
+    prompt = mk_prompt(age_band, genre, difficulty)
+
+    # PASS 1: deterministic
+    text = _gen_deterministic(MODEL, TOK, prompt)
     js = extract_json_str(text)
     obj = try_parse_json(js) if js else None
 
-    # PASS 2: sampled retry if too short or not valid JSON
+    # PASS 2: sampled retry if too short or invalid
     if not obj or len(js or "") < 40:
-        out = MODEL.generate(
-            **inputs,
-            max_new_tokens=220,
-            min_new_tokens=80,
-            do_sample=True,
-            temperature=0.6,
-            top_p=0.9,
-            top_k=50,
-            no_repeat_ngram_size=3,
-            early_stopping=True,
-        )
-        text = TOK.decode(out[0], skip_special_tokens=True)
+        text = _gen_sampled(MODEL, TOK, prompt)
         js = extract_json_str(text)
         obj = try_parse_json(js) if js else None
 
@@ -196,7 +226,7 @@ def generate(age_band, genre, difficulty):
         missing = sorted(list(need - set(obj.keys())))
         return f"ERROR: JSON missing keys {missing}\n\nPARSED:\n{json.dumps(obj, indent=2)}"
 
-    # Guardrails: exactly 4 choices and valid 0-3 answer index
+    # Guardrails
     if not isinstance(obj.get("choices"), list) or len(obj["choices"]) != 4:
         return f"ERROR: choices must be a list of 4.\n\nPARSED:\n{json.dumps(obj, indent=2)}"
     if not isinstance(obj.get("answer"), int) or not (0 <= obj["answer"] <= 3):
@@ -204,6 +234,44 @@ def generate(age_band, genre, difficulty):
 
     return json.dumps(obj, indent=2, ensure_ascii=False)
 
+# -------- Verification: base vs adapter A/B on the same prompt --------
+def verify_adapter(age_band, genre, difficulty):
+    if TOK is None or MODEL is None or ACTIVE_LABEL is None:
+        return "🛑 No adapter loaded."
+
+    prompt = mk_prompt(age_band, genre, difficulty)
+
+    # Base-only (fresh load, no adapter)
+    base_tok = AutoTokenizer.from_pretrained(ACTIVE_BASE)
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(ACTIVE_BASE)
+    base_model.eval()
+
+    base_text = _gen_deterministic(base_model, base_tok, prompt)
+    base_json = extract_json_str(base_text)
+    base_ok = bool(try_parse_json(base_json) if base_json else None)
+
+    # Adapter (current MODEL/TOK)
+    adapter_text = _gen_deterministic(MODEL, TOK, prompt)
+    adapter_json = extract_json_str(adapter_text)
+    adapter_ok = bool(try_parse_json(adapter_json) if adapter_json else None)
+
+    sha_show = (ACTIVE_SHA[:12] + "…") if ACTIVE_SHA else "unknown"
+
+    report = {
+        "active_label": ACTIVE_LABEL,
+        "base_model": ACTIVE_BASE,
+        "adapter_sha256": ACTIVE_SHA,
+        "prompt_preview": prompt[:200] + ("…" if len(prompt) > 200 else ""),
+        "base_output_preview": (base_text[:400] + "…") if len(base_text) > 400 else base_text,
+        "base_json_detected": bool(base_json),
+        "base_json_parsed_ok": base_ok,
+        "adapter_output_preview": (adapter_text[:400] + "…") if len(adapter_text) > 400 else adapter_text,
+        "adapter_json_detected": bool(adapter_json),
+        "adapter_json_parsed_ok": adapter_ok,
+        "note": "If adapter_json_parsed_ok != base_json_parsed_ok, the adapter is changing behavior."
+    }
+    return json.dumps(report, indent=2, ensure_ascii=False)
+
 def reload_list():
     items = list_available_models()
     if not items:
@@ -217,10 +285,9 @@ def do_load(label):
     can_gen = TOK is not None and MODEL is not None and ACTIVE_LABEL == label
     return status, (zip_path if os.path.isfile(zip_path) else None), gr.update(interactive=can_gen)
 
-
 # --------------------- UI ---------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## Quiz AI — Train ➜ Save ➜ Download ➜ Use (No tokens)")
+    gr.Markdown("## Quiz AI — Train ➜ Save ➜ Download ➜ Use (No tokens) + Verify Adapter")
 
     with gr.Tab("Train"):
         with gr.Row():
@@ -255,9 +322,18 @@ with gr.Blocks() as demo:
         gen_btn = gr.Button("Generate sample", interactive=False)  # disabled until model loaded
         out     = gr.Code(label="Model output (JSON expected)")
 
+        # Verify section
+        with gr.Accordion("Verify adapter (A/B vs base)", open=False):
+            v_age  = gr.Textbox("13-17", label="Age band")
+            v_gen  = gr.Dropdown(choices=["geography","science","history","math"], value="science", label="Genre")
+            v_diff = gr.Dropdown(choices=["easy","medium","hard"], value="easy", label="Difficulty")
+            verify_btn = gr.Button("Run verification")
+            verify_out = gr.Code(label="Verification report (JSON)")
+
         refresh_btn.click(fn=reload_list, outputs=[model_list, status_md, label])
         load_btn.click(fn=do_load, inputs=[model_list], outputs=[status_md, current_zip, gen_btn])
         gen_btn.click(fn=generate, inputs=[age, gen, diff], outputs=out)
+        verify_btn.click(fn=verify_adapter, inputs=[v_age, v_gen, v_diff], outputs=verify_out)
 
 if __name__ == "__main__":
     demo.launch()