Update app.py

app.py

@@ -1,206 +1,199 @@
-import gradio as gr
-import re
-import torch
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
-
-# ---------------------------
-# Load models
-# ---------------------------
-print("Loading models...")
-
-# Sentence similarity model
-retriever = pipeline("feature-extraction", model="sentence-transformers/all-MiniLM-L6-v2")
+# ─────────────────────────────────────────────────────────────────────────────
+# TruthLens – Misinformation-Aware RAG (Lite/Full modes)
+# ─────────────────────────────────────────────────────────────────────────────
 
-# NLI model (for faithfulness scoring)
-nli = pipeline("text-classification", model="facebook/bart-large-mnli", top_k=None)
+import numpy as np
+import pandas as pd
+from sklearn.metrics.pairwise import cosine_similarity
 
-# Toxicity classifier
-toxicity = pipeline("text-classification", model="unitary/toxic-bert", top_k=None)
-
-# Summarization/fallback model
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-
-# Generation model (for synthetic answers)
-M = None
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from sentence_transformers import SentenceTransformer
+
+# ===== Config =====
+GEN_TINY = "google/flan-t5-small"          # Lite mode
+GEN_FULL = "google/flan-t5-base"           # Full mode
+EMB_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+NLI_MODEL = "cross-encoder/nli-deberta-v3-small"  # light NLI (still optional)
+NER_MODEL = "dslim/bert-base-NER"
+TOX_MODEL = "unitary/toxic-bert"
+
+SAMPLE_DOCS = [
+    {"title": "WHO on Vaccines & Safety",
+     "text": "Vaccines undergo rigorous testing and continuous safety monitoring. Severe adverse reactions are rare.",
+     "url": "https://example.org/who-vaccines"},
+    {"title": "IPCC on Climate Change",
+     "text": "It is unequivocal that human influence has warmed the atmosphere, ocean and land.",
+     "url": "https://example.org/ipcc"},
+    {"title": "Elections Security Myths",
+     "text": "Independent audits reduce fraud risk; no credible evidence for nationwide manipulation.",
+     "url": "https://example.org/election-security"},
+]
+
+# ===== Lazy model holders =====
+class M:
+    emb = None
+    tok = None
+    gen = None
+    nli = None
+    ner = None
+    tox = None
+
+INDEX = {"emb": None, "texts": [], "titles": [], "urls": []}
+
+def ensure_emb():
+    if M.emb is None:
+        M.emb = SentenceTransformer(EMB_MODEL)
 
 def ensure_gen(use_tiny=True):
-    global M
-    if M is None:
-        if use_tiny:
-            model_name = "google/flan-t5-small"
-            M = pipeline("text2text-generation", model=model_name)
-        else:
-            model_name = "google/flan-t5-base"
-            M = pipeline("text2text-generation", model=model_name)
-    return M
-
-# ---------------------------
-# Demo Index (sources)
-# ---------------------------
-INDEX = {
-    "titles": [
-        "IPCC on Climate Change",
-        "Elections Security Myths",
-        "WHO on Vaccines & Safety",
-    ],
-    "urls": [
-        "https://example.org/ipcc",
-        "https://example.org/election-security",
-        "https://example.org/who-vaccines",
-    ],
-    "texts": [
-        "The IPCC states with high confidence that human activities are the dominant cause of global warming since the mid-20th century.",
-        "Studies show that widespread voter fraud in modern elections is extremely rare and not supported by credible evidence.",
-        "The World Health Organization confirms vaccines are safe and effective, with benefits vastly outweighing risks.",
-    ],
-}
-
-# ---------------------------
-# Helpers
-# ---------------------------
-def _clean(txt: str) -> str:
-    return re.sub(r"\s+", " ", txt).strip()
-
-def _extractive_fallback(question, idxs):
-    """If generation fails, fallback to extractive summarizer."""
-    ctx = " ".join([INDEX["texts"][i] for i in idxs])
-    summary = summarizer(ctx, max_length=80, min_length=30, do_sample=False)
-    return summary[0]["summary_text"], [INDEX["titles"][i] for i in idxs]
-
-# ---------------------------
-# Faithfulness scoring
-# ---------------------------
-def faithfulness_scores(answer, idxs):
-    scores, per_source = [], []
-    for i, idx in enumerate(idxs):
-        premise = INDEX["texts"][idx]
-        result = nli({"premise": premise, "hypothesis": answer})[0]
-
-        entail_score = 0.0
-        for item in result:
-            if item["label"].upper().startswith("ENTAIL"):
-                entail_score = item["score"]
-        scores.append(entail_score)
-        per_source.append((INDEX["titles"][idx], entail_score))
-
-    mean_score = sum(scores) / len(scores) if scores else 0.0
-    return mean_score, per_source
-
-# ---------------------------
-# Toxicity scoring
-# ---------------------------
-def toxicity_risk(answer):
-    result = toxicity(answer)[0]
-    toxic_score = 0.0
-    for item in result:
-        if "toxic" in item["label"].lower():
-            toxic_score = item["score"]
-    return toxic_score
-
-# ---------------------------
-# Answer generation
-# ---------------------------
-def generate_answer(question, idxs, use_tiny=True, max_new=220):
+    model_id = GEN_TINY if use_tiny else GEN_FULL
+    if (M.gen is None) or (getattr(M.gen, "_id", None) != model_id):
+        M.tok = AutoTokenizer.from_pretrained(model_id)
+        M.gen = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+        M.gen._id = model_id  # remember which is loaded
+
+def ensure_nli():
+    if M.nli is None:
+        # NOTE: no return_all_scores; we’ll use top_k=None at call time
+        M.nli = pipeline("text-classification", model=NLI_MODEL)
+
+def ensure_ner():
+    if M.ner is None:
+        M.ner = pipeline("token-classification", model=NER_MODEL, aggregation_strategy="simple")
+
+def ensure_tox():
+    if M.tox is None:
+        M.tox = pipeline("text-classification", model=TOX_MODEL)
+
+# ===== Index =====
+def build_index(extra=None):
+    ensure_emb()
+    texts = [d["text"] for d in SAMPLE_DOCS]
+    titles = [d["title"] for d in SAMPLE_DOCS]
+    urls = [d["url"] for d in SAMPLE_DOCS]
+    if extra:
+        for i, t in enumerate(extra):
+            if t and str(t).strip():
+                texts.append(str(t).strip()); titles.append(f"UserDoc {i+1}"); urls.append("user://paste")
+    INDEX["emb"] = M.emb.encode(texts, normalize_embeddings=True, convert_to_numpy=True)
+    INDEX["texts"], INDEX["titles"], INDEX["urls"] = texts, titles, urls
+
+# ===== Core steps =====
+def retrieve(q, k=3):
+    ensure_emb()
+    if INDEX["emb"] is None:
+        build_index()
+    qv = M.emb.encode([q], normalize_embeddings=True, convert_to_numpy=True)
+    sims = cosine_similarity(qv, INDEX["emb"])[0]
+    return list(np.argsort(-sims)[:k])
+
+def generate_answer(question, idxs, use_tiny=True, max_new=256):
     ensure_gen(use_tiny)
-
     ctx, cites = [], []
     for i, idx in enumerate(idxs):
         ctx.append(f"[{i+1}] {INDEX['texts'][idx]}")
         cites.append(f"[{i+1}] {INDEX['titles'][idx]} – {INDEX['urls'][idx]}")
-
-    instr = (
-        "Write a clear paragraph (3–6 sentences) that answers the user's claim "
-        "STRICTLY using the sources below. Include citations like [1], [2]. "
-        "Do not reply with only citation markers; write complete sentences."
-    )
-
-    # ✅ build ctx_block outside the f-string
-    ctx_block = "\n".join(ctx)
-
     prompt = (
-        f"{instr}\n\nSources:\n{ctx_block}\n\n"
-        f"Claim: {question}\nAnswer:"
-    )
-
-    toks = M.tokenizer(prompt, return_tensors="pt", truncation=True)
-    out = M.model.generate(
-        **toks,
-        max_new_tokens=max_new,
-        min_new_tokens=80,
-        do_sample=True,
-        temperature=0.8,
-        top_p=0.92,
-        repetition_penalty=1.15,
-        no_repeat_ngram_size=3,
-        early_stopping=True,
+        "Answer the user's claim STRICTLY using the sources below. "
+        "Use citations like [1], [2]. If unsure, say you are uncertain.\n\n"
+        f"Sources:\n{'\n'.join(ctx)}\n\n"
+        f"Question: {question}\nAnswer:"
     )
-    text = _clean(M.tokenizer.decode(out[0], skip_special_tokens=True))
-
-    if len(text) < 60 or re.fullmatch(r"\[+\d+\]+\.?", text):
-        text, cites = _extractive_fallback(question, idxs)
-
-    return text, cites
-
-# ---------------------------
-# Pipeline
-# ---------------------------
-def run_pipeline(claim, src1, src2, src3, use_tiny=True):
-    # Gather candidate sources
-    candidates = list(range(len(INDEX["texts"])))
-    if src1:
-        INDEX["texts"].append(src1)
-        INDEX["titles"].append("Custom Source 1")
-        INDEX["urls"].append("user://source1")
-        candidates.append(len(INDEX["texts"]) - 1)
-    if src2:
-        INDEX["texts"].append(src2)
-        INDEX["titles"].append("Custom Source 2")
-        INDEX["urls"].append("user://source2")
-        candidates.append(len(INDEX["texts"]) - 1)
-    if src3:
-        INDEX["texts"].append(src3)
-        INDEX["titles"].append("Custom Source 3")
-        INDEX["urls"].append("user://source3")
-        candidates.append(len(INDEX["texts"]) - 1)
-
-    # Pick top-3 sources (simplified: first 3 candidates)
-    idxs = candidates[:3]
-
-    # Generate answer
-    answer, cites = generate_answer(claim, idxs, use_tiny=use_tiny)
-
-    # Faithfulness + toxicity
-    faith_total, per_src = faithfulness_scores(answer, idxs)
-    tox = toxicity_risk(answer)
-
-    # PII redaction
-    redacted = re.sub(r"\b[A-Z][a-z]+ [A-Z][a-z]+\b", "[REDACTED]", answer)
-
-    return (
-        f"Faithfulness (mean entailment): {faith_total:.2f} | Toxicity risk: {tox:.2f}\n\n{answer}\n\n"
-        + "\n".join(cites),
-        per_src,
-        redacted,
-    )
-
-# ---------------------------
-# Gradio UI
-# ---------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## 🎯 TruthLens – Misinformation-Aware RAG\nType a claim or question and see fact-checked answers with citations.")
-
+    toks = M.tok(prompt, return_tensors="pt", truncation=True)
+    out = M.gen.generate(**toks, max_new_tokens=max_new, do_sample=False)
+    return M.tok.decode(out[0], skip_special_tokens=True), cites
+
+def nli_faithfulness(answer, idxs):
+    try:
+        ensure_nli()
+        per_src = []
+        for idx in idxs:
+            prem = INDEX["texts"][idx]
+            out = M.nli({"text": prem, "text_pair": answer}, top_k=None)
+            # Normalize shapes: out -> list -> list[dict] or dict
+            scores_obj = out[0] if isinstance(out, list) and out else out
+            scores = [scores_obj] if isinstance(scores_obj, dict) else (scores_obj or [])
+            ent = 0.0
+            for item in scores:
+                if str(item.get("label", "")).upper().startswith("ENTAIL"):
+                    ent = float(item.get("score", 0.0)); break
+            per_src.append((INDEX["titles"][idx], ent))
+        mean_ent = float(np.mean([s for _, s in per_src])) if per_src else 0.0
+        return mean_ent, per_src, None
+    except Exception as e:
+        return 0.0, [(INDEX["titles"][i], 0.0) for i in idxs], f"NLI skipped: {e}"
+
+def redact_pii(text):
+    try:
+        ensure_ner()
+        ents = M.ner(text)
+        ents = sorted(ents, key=lambda e: e.get("end",0)-e.get("start",0), reverse=True)
+        out = text
+        for e in ents:
+            s, e2 = int(e.get("start",0)), int(e.get("end",0))
+            span = text[s:e2]
+            if span:
+                out = out.replace(span, f"<{e.get('entity_group','ENT')}>")
+        return out, None
+    except Exception as e:
+        return text, f"PII redaction skipped: {e}"
+
+def tox_score(text):
+    try:
+        ensure_tox()
+        pred = M.tox(text)[0]
+        return float(pred.get("score", 0.0)), None
+    except Exception as e:
+        return 0.0, f"Toxicity check skipped: {e}"
+
+# ===== Pipeline (Lite vs Full) =====
+def run_pipeline(claim, s1, s2, s3, lite_mode):
+    # Build/refresh index with user sources
+    build_index([s1, s2, s3])
+
+    # 1) Retrieve + Generate (always on)
+    idxs = retrieve(claim, k=3)
+    answer, cites = generate_answer(claim, idxs, use_tiny=lite_mode)
+
+    # 2) Optional checks (only in Full mode, but fail-soft)
+    notes = []
+    if not lite_mode:
+        mean_ent, per_src, nli_note = nli_faithfulness(answer, idxs)
+        if nli_note: notes.append(nli_note)
+        pii, pii_note = redact_pii(answer);        redacted = pii
+        if pii_note: notes.append(pii_note)
+        tox, tox_note = tox_score(answer)
+        if tox_note: notes.append(tox_note)
+    else:
+        mean_ent, per_src = 0.0, [(INDEX["titles"][i], 0.0) for i in idxs]
+        redacted, tox = answer, 0.0
+        notes.append("Lite mode: NLI/PII/Toxicity disabled for reliability on free CPU.")
+
+    table = pd.DataFrame({"Source": [s for s,_ in per_src],
+                          "Faithfulness": [round(float(sc),3) for _, sc in per_src]})
+    summary = f"Faithfulness (mean entailment): {mean_ent:.2f} | Toxicity risk: {tox:.2f}"
+    if notes:
+        summary += "  \n" + "  \n".join(f"• {n}" for n in notes)
+    return summary, answer, "\n".join(cites), table, redacted
+
+# ===== UI =====
+with gr.Blocks(title="TruthLens – Misinformation-Aware RAG") as demo:
+    gr.Markdown("# 🧭 TruthLens – Misinformation-Aware RAG\nType a claim or question and get a grounded answer with citations.")
     with gr.Row():
-        claim = gr.Textbox(label="Claim or question", placeholder="e.g., Did humans cause climate change?")
-        src1 = gr.Textbox(label="Optional source 1", lines=3)
-        src2 = gr.Textbox(label="Optional source 2", lines=3)
-        src3 = gr.Textbox(label="Optional source 3", lines=3)
-
-    run_btn = gr.Button("Run TruthLens", variant="primary")
-
-    out_answer = gr.Textbox(label="Fact-checked answer", lines=8)
-    out_table = gr.Dataframe(headers=["Source", "Faithfulness"], label="Per-source faithfulness", wrap=True)
-    out_redact = gr.Textbox(label="PII-redacted answer", lines=6)
-
-    run_btn.click(fn=run_pipeline, inputs=[claim, src1, src2, src3], outputs=[out_answer, out_table, out_redact])
-
-demo.launch()
+        with gr.Column():
+            claim = gr.Textbox(label="Claim or question", lines=2, placeholder="e.g., Did humans cause climate change?")
+            lite = gr.Checkbox(value=True, label="Lite mode (more reliable on free CPU)")
+            run_btn = gr.Button("Run TruthLens", variant="primary")
+        with gr.Column():
+            s1 = gr.Textbox(label="Optional source 1", lines=3)
+            s2 = gr.Textbox(label="Optional source 2", lines=3)
+            s3 = gr.Textbox(label="Optional source 3", lines=3)
+    summary = gr.Markdown()
+    answer = gr.Markdown(label="Answer")
+    cites = gr.Markdown(label="Citations")
+    table = gr.Dataframe(label="Per-source faithfulness")
+    redacted = gr.Textbox(label="PII-redacted answer", lines=3)
+    run_btn.click(run_pipeline, [claim, s1, s2, s3, lite], [summary, answer, cites, table, redacted])
+
+if __name__ == "__main__":
+    demo.launch()