michlea/HallucinationDetection

1"""2Gradio demo: Detecting LLM Hallucinations from Hidden States3 4Probes Qwen2.5-0.5B internal representations to classify a response as5truthful or hallucinated — no external fact-checker, no sampling.6 7Usage (HF Spaces or local):8    python app.py9 10Requires probe.joblib produced by save_probe.py.11"""12from __future__ import annotations13 14import os15import re16import warnings17 18import gradio as gr19import joblib20import numpy as np21import torch22from transformers import AutoModelForCausalLM, AutoTokenizer23 24# ── Silence noisy HF / torch logs ─────────────────────────────────────────────25os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")26warnings.filterwarnings("ignore")27 28# ── Config ─────────────────────────────────────────────────────────────────────29MODEL_NAME = "Qwen/Qwen2.5-0.5B"30MAX_LENGTH = 51231DEVICE = "cuda" if torch.cuda.is_available() else "cpu"32DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float3233 34# Qwen2.5-0.5B: 24 transformer layers + embedding = 25 hidden states.35# Mid-band (40–64% depth) = layers 9–16; signal peaks here, not at the final layer.36MID_BAND = list(range(9, 17))37GEO_LAYERS = (8, 12, 16, 20, 24)38 39_WORD = re.compile(r"\w+")40 41# ── Load model once at startup ─────────────────────────────────────────────────42print(f"[startup] Loading {MODEL_NAME} on {DEVICE} ({DTYPE})…")43tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)44if tokenizer.pad_token is None:45    tokenizer.pad_token = tokenizer.eos_token46model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=DTYPE)47model.eval()48model.to(DEVICE)49print("[startup] Model loaded.")50 51# ── Load probe ────────────────────────────────────────────────────────────────52_artifact = joblib.load("probe.joblib")53_probe = _artifact["probe"]54_threshold = float(_artifact["threshold"])55print(f"[startup] Probe loaded  experiment={_artifact.get('exp_name')}  "56      f"threshold={_threshold:.3f}  features={_artifact.get('n_features')}")57 58 59# ── Prompt formatting ──────────────────────────────────────────────────────────60def _fmt_prompt(context: str, question: str) -> str:61    """Reconstruct the ChatML prompt used during Qwen SMILES data generation."""62    return (63        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"64        "<|im_start|>user\n"65        "Given the context, answer the question in a single brief but complete sentence.\n"66        f"{context}\n"67        "Note that your answer must be based only on the context. "68        'If the context does not provide enough information, '69        'reply with: "Unable to answer based on given context".\n'70        f"Here is the question: {question}\nYour answer:<|im_end|>\n"71        "<|im_start|>assistant\n"72    )73 74 75# ── Inference helpers ──────────────────────────────────────────────────────────76@torch.no_grad()77def _generate(prompt: str) -> str:78    """Let Qwen2.5-0.5B produce one greedy response."""79    enc = tokenizer(80        prompt, return_tensors="pt",81        truncation=True, max_length=MAX_LENGTH - 80,82    ).to(DEVICE)83    out = model.generate(84        **enc,85        max_new_tokens=80,86        do_sample=False,87        pad_token_id=tokenizer.eos_token_id,88    )89    new_ids = out[0][enc.input_ids.shape[1]:]90    return tokenizer.decode(new_ids, skip_special_tokens=True).strip()91 92 93@torch.no_grad()94def _extract(text: str, prompt_len: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:95    """Run one forward pass and return mid-band pooled + geometric features."""96    enc = tokenizer(97        text, return_tensors="pt",98        truncation=True, max_length=MAX_LENGTH,99    ).to(DEVICE)100 101    out = model(**enc, output_hidden_states=True)102    hs = out.hidden_states  # 25 × (1, T, 896)103 104    am = enc.attention_mask[0].bool()105    n_real = int(am.sum())106 107    # Stack all layers, real tokens only: (25, n_real, 896) float32108    h = torch.stack([layer[0][am] for layer in hs], dim=0).float()109 110    rfrom = max(min(prompt_len, n_real - 1), 0)111    resp = h[:, rfrom:, :] if rfrom < h.size(1) else h[:, -1:, :]112 113    # ── Pooled hidden states: max + std over mid-band layers ──────────────────114    h_max = resp[MID_BAND].amax(dim=1).cpu().numpy().astype(np.float32)   # (8, 896)115    h_std = (116        resp[MID_BAND].std(dim=1).cpu().numpy().astype(np.float32)117        if resp.size(1) > 1118        else np.zeros((len(MID_BAND), h.shape[-1]), dtype=np.float32)119    )120 121    # ── Geometric scalars (EigenScore family) ─────────────────────────────────122    geo_vals: list[float] = []123    prev_mean: torch.Tensor | None = None124    for L in GEO_LAYERS:125        layer = resp[L]                                   # (n_resp, 896)126        norm = layer.norm(dim=-1).mean().item()127 128        centered = layer - layer.mean(dim=0, keepdim=True)129        n = max(layer.size(0), 1)130        gram = centered @ centered.T / n131        gram = gram + 1e-4 * torch.eye(gram.size(0), device=gram.device, dtype=gram.dtype)132        try:133            eig = torch.linalg.eigvalsh(gram)134        except Exception:135            eig = torch.linalg.eigvalsh(gram.double().cpu()).float()136        eigenscore = torch.log(eig.clamp_min(1e-6)).mean().item()137 138        cur_mean = layer.mean(dim=0)139        drift = (140            torch.nn.functional.cosine_similarity(141                cur_mean.unsqueeze(0), prev_mean.unsqueeze(0),142            ).item()143            if prev_mean is not None else 1.0144        )145        prev_mean = cur_mean146        geo_vals += [norm, eigenscore, drift]147 148    return h_max, h_std, np.array(geo_vals, dtype=np.float32)149 150 151def _surface(context: str, response: str) -> np.ndarray:152    """Cheap lexical grounding + length features (6 scalars)."""153    ctx_words = set(_WORD.findall(context.lower()))154    resp_words = _WORD.findall(response.lower())155    n = max(len(resp_words), 1)156    grounding = sum(w in ctx_words for w in resp_words) / n157    return np.array([158        grounding,159        1.0 - grounding,160        float(len(resp_words)),161        float(len(set(resp_words)) / n),162        float("unable to answer" in response.lower()),163        float("unable to answer" in context.lower()),164    ], dtype=np.float32)165 166 167# ── Main detection function ────────────────────────────────────────────────────168def detect(169    context: str,170    question: str,171    response: str,172    auto_gen: bool,173    progress=gr.Progress(track_tqdm=True),174):175    context, question = context.strip(), question.strip()176 177    if not context:178        return "", 0.0, "⚠️ Please provide a context passage.", ""179    if not question:180        return "", 0.0, "⚠️ Please provide a question.", ""181 182    prompt = _fmt_prompt(context, question)183 184    # ── Step 1: get response ───────────────────────────────────────────────────185    if auto_gen or not response.strip():186        progress(0.05, desc="Generating response with Qwen2.5-0.5B…")187        response = _generate(prompt)188 189    # ── Step 2: extract features ───────────────────────────────────────────────190    progress(0.3, desc="Running forward pass (hidden states)…")191    text = prompt + response192    prompt_len = len(tokenizer(prompt, truncation=True, max_length=MAX_LENGTH)["input_ids"])193    h_max, h_std, geo = _extract(text, prompt_len)194    surf = _surface(context, response)195 196    # Feature vector must match build_matrix order:197    #   [max_resp_L9..L16 flat (7168), std_resp_L9..L16 flat (7168), geo (15), surf (6)]198    X = np.concatenate([h_max.flatten(), h_std.flatten(), geo, surf]).reshape(1, -1)199 200    # ── Step 3: probe ──────────────────────────────────────────────────────────201    progress(0.9, desc="Running linear probe…")202    prob = float(_probe.predict_proba(X)[0, 1])203    grounding = float(surf[0])204    resp_len = int(surf[2])205 206    # ── Step 4: format output ──────────────────────────────────────────────────207    is_hallu = prob >= _threshold208    verdict = "🔴  HALLUCINATED" if is_hallu else "🟢  TRUTHFUL"209 210    bar = "█" * round(prob * 20) + "░" * (20 - round(prob * 20))211    detail = (212        f"**Probe score:** `{prob:.3f}` / threshold `{_threshold:.3f}`\n\n"213        f"`[{bar}]` {prob:.1%} hallucination probability\n\n"214        f"| Signal | Value |\n"215        f"|---|---|\n"216        f"| Lexical grounding (resp words in context) | {grounding:.1%} |\n"217        f"| Response length | {resp_len} words |\n"218        f"| Contains 'unable to answer' | {'yes' if surf[4] else 'no'} |\n\n"219        f"*Hidden-state features: {len(MID_BAND)} mid-stack layers (9–16 of 25), "220        f"`max` + `std` pooling over response tokens.*"221    )222    223    progress(1.0)224    return response, round(prob, 4), verdict, detail225 226 227# ── UI ─────────────────────────────────────────────────────────────────────────228_TITLE = "# Hallucination Detection via Hidden-State Probing"229_DESC = """\230Detects whether a Qwen2.5-0.5B response is **hallucinated** or **truthful** by reading the231model's own internal representations — **no external checker, no sampling**.232 233**Key finding from the research:** the truthfulness signal peaks at *middle layers* (~40–60%234depth), not the final layer. A simple linear probe on mid-stack hidden states reaches235**AUROC ≈ 0.77**, vs. 0.67 for the standard last-layer baseline.236 237*Trained on SMILES: 689 SQuAD-derived QA samples answered by Qwen2.5-0.5B.*238"""239 240_HOW = """\241---242**How it works**243 2441. Your `context + question + response` is tokenised and fed through Qwen2.5-0.5B in a245   *single forward pass* (no generation needed if you supply the response).2462. Hidden states from 8 mid-stack layers (9–16) are pooled over the *response tokens* using247   `max` and `std` pooling — the two operations that carry the strongest truthfulness signal.2483. EigenScore-style geometric scalars (log-det of the response-token covariance) and six249   cheap lexical features (grounding ratio, length, …) are appended.2504. A strong-L2 logistic regression (`fusion_nopca_L2`) returns a hallucination probability.251   A nested-CV-tuned threshold converts it to a binary verdict.252"""253 254EXAMPLES = [255    [256        "The United States Geological Survey (USGS) released a new earthquake forecast for California, "257        "predicting a 60 percent probability of a major earthquake in the next 30 years.",258        "Which organization released a California earthquake forecast?",259        "Here is the answer: USGS",260        False,261    ],262    [263        "A ring is an algebraic structure in which addition and multiplication are defined and have "264        "similar properties to those operations defined for the integers. Ring theory is the study of rings.",265        "What is the name of an algebraic structure in which addition, subtraction and multiplication are defined?",266        "Prime number",267        False,268    ],269    [270        "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. "271        "It is named after the engineer Gustave Eiffel, whose company designed and built the tower "272        "from 1887 to 1889.",273        "When was the Eiffel Tower built?",274        "",275        True,276    ],277    [278        "Marie Curie was a Polish and naturalised-French physicist and chemist who conducted pioneering "279        "research on radioactivity. She was the first woman to win a Nobel Prize, the first person to "280        "win the Nobel Prize twice, and the only person to win the Nobel Prize in two scientific sciences.",281        "How many Nobel Prizes did Marie Curie win?",282        "Marie Curie won the Nobel Prize three times.",283        False,284    ],285]286 287with gr.Blocks(title="Hallucination Detector", theme=gr.themes.Soft()) as demo:288    gr.Markdown(_TITLE)289    gr.Markdown(_DESC)290 291    with gr.Row():292        # ── Left: inputs ───────────────────────────────────────────────────────293        with gr.Column(scale=1):294            ctx_in = gr.Textbox(295                label="Context passage",296                lines=7,297                placeholder="Paste the context paragraph the model should answer from…",298            )299            q_in = gr.Textbox(300                label="Question",301                lines=2,302                placeholder="What question is being asked?",303            )304            resp_in = gr.Textbox(305                label="Response to evaluate  (leave blank to auto-generate)",306                lines=3,307                placeholder="The model's answer…",308            )309            auto_cb = gr.Checkbox(310                label="Auto-generate response from Qwen2.5-0.5B",311                value=False,312            )313            run_btn = gr.Button("Detect hallucination", variant="primary", size="lg")314 315        # ── Right: outputs ──────────────────────────────────────────────────────316        with gr.Column(scale=1):317            resp_out = gr.Textbox(318                label="Response evaluated",319                lines=3,320                interactive=False,321            )322            prob_out = gr.Number(label="Hallucination probability  (0 = truthful, 1 = hallucinated)")323            verdict_out = gr.Textbox(label="Verdict", interactive=False, lines=1)324            detail_md = gr.Markdown()325 326    gr.Examples(327        examples=EXAMPLES,328        inputs=[ctx_in, q_in, resp_in, auto_cb],329        outputs=[resp_out, prob_out, verdict_out, detail_md],330        fn=detect,331        cache_examples=False,332        label="Try an example  (one truthful · two hallucinated · one auto-generated)",333    )334 335    run_btn.click(336        detect,337        inputs=[ctx_in, q_in, resp_in, auto_cb],338        outputs=[resp_out, prob_out, verdict_out, detail_md],339    )340 341    gr.Markdown(_HOW)342 343demo.launch()344