app.py

"""
Gradio demo: Detecting LLM Hallucinations from Hidden States

Probes Qwen2.5-0.5B internal representations to classify a response as
truthful or hallucinated — no external fact-checker, no sampling.

Usage (HF Spaces or local):
    python app.py

Requires probe.joblib produced by save_probe.py.
"""
from __future__ import annotations

import os
import re
import warnings

import gradio as gr
import joblib
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ── Silence noisy HF / torch logs ─────────────────────────────────────────────
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
warnings.filterwarnings("ignore")

# ── Config ─────────────────────────────────────────────────────────────────────
MODEL_NAME = "Qwen/Qwen2.5-0.5B"
MAX_LENGTH = 512
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32

# Qwen2.5-0.5B: 24 transformer layers + embedding = 25 hidden states.
# Mid-band (40–64% depth) = layers 9–16; signal peaks here, not at the final layer.
MID_BAND = list(range(9, 17))
GEO_LAYERS = (8, 12, 16, 20, 24)

_WORD = re.compile(r"\w+")

# ── Load model once at startup ─────────────────────────────────────────────────
print(f"[startup] Loading {MODEL_NAME} on {DEVICE} ({DTYPE})…")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=DTYPE)
model.eval()
model.to(DEVICE)
print("[startup] Model loaded.")

# ── Load probe ────────────────────────────────────────────────────────────────
_artifact = joblib.load("probe.joblib")
_probe = _artifact["probe"]
_threshold = float(_artifact["threshold"])
print(f"[startup] Probe loaded  experiment={_artifact.get('exp_name')}  "
      f"threshold={_threshold:.3f}  features={_artifact.get('n_features')}")


# ── Prompt formatting ──────────────────────────────────────────────────────────
def _fmt_prompt(context: str, question: str) -> str:
    """Reconstruct the ChatML prompt used during Qwen SMILES data generation."""
    return (
        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
        "<|im_start|>user\n"
        "Given the context, answer the question in a single brief but complete sentence.\n"
        f"{context}\n"
        "Note that your answer must be based only on the context. "
        'If the context does not provide enough information, '
        'reply with: "Unable to answer based on given context".\n'
        f"Here is the question: {question}\nYour answer:<|im_end|>\n"
        "<|im_start|>assistant\n"
    )


# ── Inference helpers ──────────────────────────────────────────────────────────
@torch.no_grad()
def _generate(prompt: str) -> str:
    """Let Qwen2.5-0.5B produce one greedy response."""
    enc = tokenizer(
        prompt, return_tensors="pt",
        truncation=True, max_length=MAX_LENGTH - 80,
    ).to(DEVICE)
    out = model.generate(
        **enc,
        max_new_tokens=80,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )
    new_ids = out[0][enc.input_ids.shape[1]:]
    return tokenizer.decode(new_ids, skip_special_tokens=True).strip()


@torch.no_grad()
def _extract(text: str, prompt_len: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Run one forward pass and return mid-band pooled + geometric features."""
    enc = tokenizer(
        text, return_tensors="pt",
        truncation=True, max_length=MAX_LENGTH,
    ).to(DEVICE)

    out = model(**enc, output_hidden_states=True)
    hs = out.hidden_states  # 25 × (1, T, 896)

    am = enc.attention_mask[0].bool()
    n_real = int(am.sum())

    # Stack all layers, real tokens only: (25, n_real, 896) float32
    h = torch.stack([layer[0][am] for layer in hs], dim=0).float()

    rfrom = max(min(prompt_len, n_real - 1), 0)
    resp = h[:, rfrom:, :] if rfrom < h.size(1) else h[:, -1:, :]

    # ── Pooled hidden states: max + std over mid-band layers ──────────────────
    h_max = resp[MID_BAND].amax(dim=1).cpu().numpy().astype(np.float32)   # (8, 896)
    h_std = (
        resp[MID_BAND].std(dim=1).cpu().numpy().astype(np.float32)
        if resp.size(1) > 1
        else np.zeros((len(MID_BAND), h.shape[-1]), dtype=np.float32)
    )

    # ── Geometric scalars (EigenScore family) ─────────────────────────────────
    geo_vals: list[float] = []
    prev_mean: torch.Tensor | None = None
    for L in GEO_LAYERS:
        layer = resp[L]                                   # (n_resp, 896)
        norm = layer.norm(dim=-1).mean().item()

        centered = layer - layer.mean(dim=0, keepdim=True)
        n = max(layer.size(0), 1)
        gram = centered @ centered.T / n
        gram = gram + 1e-4 * torch.eye(gram.size(0), device=gram.device, dtype=gram.dtype)
        try:
            eig = torch.linalg.eigvalsh(gram)
        except Exception:
            eig = torch.linalg.eigvalsh(gram.double().cpu()).float()
        eigenscore = torch.log(eig.clamp_min(1e-6)).mean().item()

        cur_mean = layer.mean(dim=0)
        drift = (
            torch.nn.functional.cosine_similarity(
                cur_mean.unsqueeze(0), prev_mean.unsqueeze(0),
            ).item()
            if prev_mean is not None else 1.0
        )
        prev_mean = cur_mean
        geo_vals += [norm, eigenscore, drift]

    return h_max, h_std, np.array(geo_vals, dtype=np.float32)


def _surface(context: str, response: str) -> np.ndarray:
    """Cheap lexical grounding + length features (6 scalars)."""
    ctx_words = set(_WORD.findall(context.lower()))
    resp_words = _WORD.findall(response.lower())
    n = max(len(resp_words), 1)
    grounding = sum(w in ctx_words for w in resp_words) / n
    return np.array([
        grounding,
        1.0 - grounding,
        float(len(resp_words)),
        float(len(set(resp_words)) / n),
        float("unable to answer" in response.lower()),
        float("unable to answer" in context.lower()),
    ], dtype=np.float32)


# ── Main detection function ────────────────────────────────────────────────────
def detect(
    context: str,
    question: str,
    response: str,
    auto_gen: bool,
    progress=gr.Progress(track_tqdm=True),
):
    context, question = context.strip(), question.strip()

    if not context:
        return "", 0.0, "⚠️ Please provide a context passage.", ""
    if not question:
        return "", 0.0, "⚠️ Please provide a question.", ""

    prompt = _fmt_prompt(context, question)

    # ── Step 1: get response ───────────────────────────────────────────────────
    if auto_gen or not response.strip():
        progress(0.05, desc="Generating response with Qwen2.5-0.5B…")
        response = _generate(prompt)

    # ── Step 2: extract features ───────────────────────────────────────────────
    progress(0.3, desc="Running forward pass (hidden states)…")
    text = prompt + response
    prompt_len = len(tokenizer(prompt, truncation=True, max_length=MAX_LENGTH)["input_ids"])
    h_max, h_std, geo = _extract(text, prompt_len)
    surf = _surface(context, response)

    # Feature vector must match build_matrix order:
    #   [max_resp_L9..L16 flat (7168), std_resp_L9..L16 flat (7168), geo (15), surf (6)]
    X = np.concatenate([h_max.flatten(), h_std.flatten(), geo, surf]).reshape(1, -1)

    # ── Step 3: probe ──────────────────────────────────────────────────────────
    progress(0.9, desc="Running linear probe…")
    prob = float(_probe.predict_proba(X)[0, 1])
    grounding = float(surf[0])
    resp_len = int(surf[2])

    # ── Step 4: format output ──────────────────────────────────────────────────
    is_hallu = prob >= _threshold
    verdict = "🔴  HALLUCINATED" if is_hallu else "🟢  TRUTHFUL"

    bar = "█" * round(prob * 20) + "░" * (20 - round(prob * 20))
    detail = (
        f"**Probe score:** `{prob:.3f}` / threshold `{_threshold:.3f}`\n\n"
        f"`[{bar}]` {prob:.1%} hallucination probability\n\n"
        f"| Signal | Value |\n"
        f"|---|---|\n"
        f"| Lexical grounding (resp words in context) | {grounding:.1%} |\n"
        f"| Response length | {resp_len} words |\n"
        f"| Contains 'unable to answer' | {'yes' if surf[4] else 'no'} |\n\n"
        f"*Hidden-state features: {len(MID_BAND)} mid-stack layers (9–16 of 25), "
        f"`max` + `std` pooling over response tokens.*"
    )
    
    progress(1.0)
    return response, round(prob, 4), verdict, detail


# ── UI ─────────────────────────────────────────────────────────────────────────
_TITLE = "# Hallucination Detection via Hidden-State Probing"
_DESC = """\
Detects whether a Qwen2.5-0.5B response is **hallucinated** or **truthful** by reading the
model's own internal representations — **no external checker, no sampling**.

**Key finding from the research:** the truthfulness signal peaks at *middle layers* (~40–60%
depth), not the final layer. A simple linear probe on mid-stack hidden states reaches
**AUROC ≈ 0.77**, vs. 0.67 for the standard last-layer baseline.

*Trained on SMILES: 689 SQuAD-derived QA samples answered by Qwen2.5-0.5B.*
"""

_HOW = """\
---
**How it works**

1. Your `context + question + response` is tokenised and fed through Qwen2.5-0.5B in a
   *single forward pass* (no generation needed if you supply the response).
2. Hidden states from 8 mid-stack layers (9–16) are pooled over the *response tokens* using
   `max` and `std` pooling — the two operations that carry the strongest truthfulness signal.
3. EigenScore-style geometric scalars (log-det of the response-token covariance) and six
   cheap lexical features (grounding ratio, length, …) are appended.
4. A strong-L2 logistic regression (`fusion_nopca_L2`) returns a hallucination probability.
   A nested-CV-tuned threshold converts it to a binary verdict.
"""

EXAMPLES = [
    [
        "The United States Geological Survey (USGS) released a new earthquake forecast for California, "
        "predicting a 60 percent probability of a major earthquake in the next 30 years.",
        "Which organization released a California earthquake forecast?",
        "Here is the answer: USGS",
        False,
    ],
    [
        "A ring is an algebraic structure in which addition and multiplication are defined and have "
        "similar properties to those operations defined for the integers. Ring theory is the study of rings.",
        "What is the name of an algebraic structure in which addition, subtraction and multiplication are defined?",
        "Prime number",
        False,
    ],
    [
        "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. "
        "It is named after the engineer Gustave Eiffel, whose company designed and built the tower "
        "from 1887 to 1889.",
        "When was the Eiffel Tower built?",
        "",
        True,
    ],
    [
        "Marie Curie was a Polish and naturalised-French physicist and chemist who conducted pioneering "
        "research on radioactivity. She was the first woman to win a Nobel Prize, the first person to "
        "win the Nobel Prize twice, and the only person to win the Nobel Prize in two scientific sciences.",
        "How many Nobel Prizes did Marie Curie win?",
        "Marie Curie won the Nobel Prize three times.",
        False,
    ],
]

with gr.Blocks(title="Hallucination Detector", theme=gr.themes.Soft()) as demo:
    gr.Markdown(_TITLE)
    gr.Markdown(_DESC)

    with gr.Row():
        # ── Left: inputs ───────────────────────────────────────────────────────
        with gr.Column(scale=1):
            ctx_in = gr.Textbox(
                label="Context passage",
                lines=7,
                placeholder="Paste the context paragraph the model should answer from…",
            )
            q_in = gr.Textbox(
                label="Question",
                lines=2,
                placeholder="What question is being asked?",
            )
            resp_in = gr.Textbox(
                label="Response to evaluate  (leave blank to auto-generate)",
                lines=3,
                placeholder="The model's answer…",
            )
            auto_cb = gr.Checkbox(
                label="Auto-generate response from Qwen2.5-0.5B",
                value=False,
            )
            run_btn = gr.Button("Detect hallucination", variant="primary", size="lg")

        # ── Right: outputs ──────────────────────────────────────────────────────
        with gr.Column(scale=1):
            resp_out = gr.Textbox(
                label="Response evaluated",
                lines=3,
                interactive=False,
            )
            prob_out = gr.Number(label="Hallucination probability  (0 = truthful, 1 = hallucinated)")
            verdict_out = gr.Textbox(label="Verdict", interactive=False, lines=1)
            detail_md = gr.Markdown()

    gr.Examples(
        examples=EXAMPLES,
        inputs=[ctx_in, q_in, resp_in, auto_cb],
        outputs=[resp_out, prob_out, verdict_out, detail_md],
        fn=detect,
        cache_examples=False,
        label="Try an example  (one truthful · two hallucinated · one auto-generated)",
    )

    run_btn.click(
        detect,
        inputs=[ctx_in, q_in, resp_in, auto_cb],
        outputs=[resp_out, prob_out, verdict_out, detail_md],
    )

    gr.Markdown(_HOW)

demo.launch()