app.py

import os
import json
import re
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import faiss
import torch
import nltk
import gradio as gr

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from transformers import pipeline as hf_pipeline
from langchain_core.prompts import PromptTemplate
from nltk.tokenize import sent_tokenize

nltk.download("punkt", quiet=True)
nltk.download("punkt_tab", quiet=True)

#  Paths 
FAISS_PATH = "wiki_index.faiss"
META_PATH  = "wiki_meta.json"

#   Load FAISS index + chunk metadata
print("Loading FAISS index and metadata...")
index = faiss.read_index(FAISS_PATH)

with open(META_PATH, "r") as f:
    meta = json.load(f)

chunks     = meta["chunks"]
chunk_meta = meta["meta"]
print(f"Loaded {index.ntotal} vectors and {len(chunks)} chunks")

 # Load embedding model (must match the one used to build the index) 
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
print(f"Loading embedder: {EMBEDDING_MODEL}")
embedder = SentenceTransformer(EMBEDDING_MODEL)

# Load small LLM for CPU deployment 
LLM_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
print(f"Loading LLM: {LLM_MODEL}  (CPU, fp32 — this may take a minute)...")
tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
model     = AutoModelForCausalLM.from_pretrained(LLM_MODEL, torch_dtype=torch.float32)
model.eval()

hf_pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=300,
    temperature=0.1,
    do_sample=True,
    repetition_penalty=1.1,
    return_full_text=False,
)
print("LLM loaded")

# Load NLI model for hallucination detection 
print("Loading NLI model...")
nli_pipe = hf_pipeline(
    "text-classification",
    model="cross-encoder/nli-deberta-v3-small",
    device=-1,  # CPU
)
print("NLI model loaded")

# Prompt template
PROMPT_TEMPLATE = """<|im_start|>system
You are a helpful assistant. Answer the question using ONLY the context below.
If the context does not contain enough information, say exactly:
"I don't know based on the provided documents."
Do NOT use any outside knowledge.
<|im_end|>
<|im_start|>user
Context:
{context}

Question: {question}
<|im_end|>
<|im_start|>assistant
"""

prompt = PromptTemplate(
    input_variables=["context", "question"],
    template=PROMPT_TEMPLATE,
)

#  Meta-sentence filter 
META_PATTERNS = [
    r"i don'?t know based on the provided documents?",
    r"according to the (provided )?context",
    r"based on the (provided )?(context|documents?)",
    r"the context does not (provide|contain|mention)",
    r"this information (was|is) not (present|mentioned|provided)",
    r"not (mentioned|stated|provided) in the (context|documents?)",
]

def is_meta_sentence(sentence: str) -> bool:
    s = sentence.lower().strip()
    return any(re.search(p, s) for p in META_PATTERNS)

#  Retrieval 
def retrieve(query, k=3):
    q_vec = embedder.encode([query], convert_to_numpy=True).astype(np.float32)
    distances, indices = index.search(q_vec, k)
    seen    = set()
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        title = chunk_meta[idx]["title"]
        if title not in seen:          
            seen.add(title)
            results.append({
                "chunk": chunks[idx],
                "title": title,
                "score": float(dist),
            })
    return results

#  RAG query
def rag_query(question, k=3):
    retrieved = retrieve(question, k=k)
    context   = "\n\n".join([f"[{r['title']}]\n{r['chunk']}" for r in retrieved])
    filled    = prompt.format(context=context, question=question)
    output    = hf_pipe(filled)
    answer    = output[0]["generated_text"].strip()
    return {"question": question, "answer": answer,
            "context": context, "retrieved": retrieved}

#  Hallucination detection  
def detect_hallucination(answer, context, threshold=0.3):
    sentences      = sent_tokenize(answer)
    context_chunks = [c.strip() for c in context.split("\n\n") if len(c.strip()) > 20]
    results        = []

    for sent in sentences:
        sent = sent.strip()
        if len(sent) < 10:
            continue

        # Skip meta/hedge sentences 
        if is_meta_sentence(sent):
            results.append({
                "sentence":     sent,
                "label":        "META",
                "entail_score": None,
            })
            continue

        best_score = 0.0
        for chunk in context_chunks:
            
            nli_out = nli_pipe(
                {"text": chunk, "text_pair": sent},
                truncation=True,
                max_length=512,
                top_k=None,
            )
            for item in nli_out:
                if item["label"].lower() == "entailment":
                    best_score = max(best_score, item["score"])

        results.append({
            "sentence":     sent,
            "label":        "GROUNDED" if best_score >= threshold else "HALLUCINATED",
            "entail_score": round(best_score, 3),
        })

    return results

# Session stats 
session_stats = {"total": 0, "grounded": 0, "hallucinated": 0}

def build_stats_html():
    t   = session_stats["total"]
    g   = session_stats["grounded"]
    h   = session_stats["hallucinated"]
    pct = int(100 * g / max(t, 1))
    return f"""
    <div style='display:flex;justify-content:space-around;padding:12px;
                background:#1a1d2e;border-radius:12px;border:1px solid #2a2d3e'>
      <div style='text-align:center'>
        <div style='font-size:1.5em;font-weight:700;color:#4f8ef7'>{t}</div>
        <div style='font-size:0.72em;color:#8b8fa8;text-transform:uppercase'>Sentences</div>
      </div>
      <div style='text-align:center'>
        <div style='font-size:1.5em;font-weight:700;color:#2ecc71'>{g}</div>
        <div style='font-size:0.72em;color:#8b8fa8;text-transform:uppercase'>Grounded</div>
      </div>
      <div style='text-align:center'>
        <div style='font-size:1.5em;font-weight:700;color:#e74c3c'>{h}</div>
        <div style='font-size:0.72em;color:#8b8fa8;text-transform:uppercase'>Hallucinated</div>
      </div>
      <div style='text-align:center'>
        <div style='font-size:1.5em;font-weight:700;color:#f39c12'>{pct}%</div>
        <div style='font-size:0.72em;color:#8b8fa8;text-transform:uppercase'>Faithfulness</div>
      </div>
    </div>
    """

# Main demo function
def truthrag_demo(question, num_chunks):
    if not question.strip():
        return (
            "<div style='color:#8b8fa8;text-align:center;padding:40px'>"
            "Enter a question above and click <strong>Ask TruthRAG</strong></div>",
            "",
            build_stats_html(),
        )

    result     = rag_query(question, k=int(num_chunks))
    detections = detect_hallucination(result["answer"], result["context"])

    # Update session stats 
    for d in detections:
        if d["label"] == "META":
            continue
        session_stats["total"] += 1
        if d["label"] == "GROUNDED":
            session_stats["grounded"] += 1
        else:
            session_stats["hallucinated"] += 1

    scored     = [d for d in detections if d["label"] != "META"]
    grounded   = sum(1 for d in scored if d["label"] == "GROUNDED")
    total      = len(scored)
    faith_pct  = int(100 * grounded / max(total, 1))

    if faith_pct >= 70:
        faith_color, faith_label = "#2ecc71", "High"
    elif faith_pct >= 40:
        faith_color, faith_label = "#f39c12", "Medium"
    else:
        faith_color, faith_label = "#e74c3c", "Low"

    html = f"""
    <div style='font-family:Segoe UI,sans-serif;padding:4px'>
      <div style='display:flex;align-items:center;gap:10px;margin-bottom:16px;
                  padding:10px 14px;background:#0f1117;border-radius:10px;
                  border:1px solid #2a2d3e'>
        <span style='font-size:0.8em;color:#8b8fa8;text-transform:uppercase;
                     letter-spacing:0.4px'>Faithfulness</span>
        <span style='font-size:1.1em;font-weight:700;color:{faith_color}'>
          {faith_pct}% — {faith_label}
        </span>
        <span style='margin-left:auto;font-size:0.8em;color:#8b8fa8'>
          {grounded}/{total} sentences grounded
        </span>
      </div>
      <div style='display:flex;flex-direction:column;gap:8px'>
    """

    for d in detections:
        if d["label"] == "META":
            # Show meta sentences in a neutral style, not scored
            html += f"""
            <div style='background:rgba(139,143,168,0.08);border-left:3px solid #8b8fa8;
                        border-radius:8px;padding:12px 14px'>
              <span style='font-size:0.72em;color:#8b8fa8;font-style:italic'>
                [system response — not scored]
              </span>
              <p style='color:#8b8fa8;font-size:0.92em;line-height:1.6;margin:4px 0 0'>
                {d['sentence']}
              </p>
            </div>
            """
        elif d["label"] == "GROUNDED":
            html += f"""
            <div style='background:rgba(46,204,113,0.08);border-left:3px solid #2ecc71;
                        border-radius:8px;padding:12px 14px'>
              <div style='display:flex;justify-content:space-between;
                          align-items:center;margin-bottom:6px'>
                <span style='background:rgba(46,204,113,0.15);color:#2ecc71;
                             font-size:0.72em;font-weight:600;padding:2px 8px;
                             border-radius:20px;text-transform:uppercase'>GROUNDED</span>
                <span style='color:#8b8fa8;font-size:0.75em'>
                  score: {d['entail_score']}
                </span>
              </div>
              <p style='color:#e0e3ef;font-size:0.92em;line-height:1.6;margin:0'>
                {d['sentence']}
              </p>
            </div>
            """
        else:
            html += f"""
            <div style='background:rgba(231,76,60,0.08);border-left:3px solid #e74c3c;
                        border-radius:8px;padding:12px 14px'>
              <div style='display:flex;justify-content:space-between;
                          align-items:center;margin-bottom:6px'>
                <span style='background:rgba(231,76,60,0.15);color:#e74c3c;
                             font-size:0.72em;font-weight:600;padding:2px 8px;
                             border-radius:20px;text-transform:uppercase'>HALLUCINATED</span>
                <span style='color:#8b8fa8;font-size:0.75em'>
                  score: {d['entail_score']}
                </span>
              </div>
              <p style='color:#e0e3ef;font-size:0.92em;line-height:1.6;margin:0'>
                {d['sentence']}
              </p>
            </div>
            """

    html += "</div></div>"

    # Sources deduplicated
    sources = ""
    for r in result["retrieved"]:
        sources += f"** {r['title']}**\n{r['chunk'][:400]}...\n\n---\n\n"

    return html, sources, build_stats_html()


# Gradio UI
custom_css = """
* { box-sizing: border-box; }
body, .gradio-container {
    background-color: #0f1117 !important;
    font-family: 'Segoe UI', sans-serif !important;
}
textarea, input[type="text"] {
    background: #0f1117 !important;
    border: 1px solid #2a2d3e !important;
    color: #ffffff !important;
    border-radius: 10px !important;
}
textarea:focus, input:focus {
    border-color: #4f8ef7 !important;
    box-shadow: 0 0 0 2px rgba(79,142,247,0.15) !important;
}
.ask-btn {
    background: linear-gradient(135deg, #4f8ef7, #7b5ea7) !important;
    color: white !important;
    border: none !important;
    border-radius: 10px !important;
    font-weight: 600 !important;
    width: 100% !important;
}
input[type="range"] { accent-color: #4f8ef7 !important; }
label, .gr-label {
    color: #8b8fa8 !important;
    font-size: 0.82em !important;
    text-transform: uppercase !important;
}
"""

with gr.Blocks(css=custom_css, title="TruthRAG") as demo:

    gr.HTML("""
    <div style='text-align:center;padding:30px 20px 10px'>
      <h1 style='font-size:2.4em;font-weight:700;color:#ffffff;margin:0'>
         TruthRAG
      </h1>
      <p style='color:#8b8fa8;font-size:1em;margin-top:8px'>
        Retrieval-Augmented Generation with Automated Hallucination Detection
      </p>
    </div>
    """)

    gr.HTML("""
    <div style='background:#1a1d2e;border:1px solid #2a2d3e;border-radius:10px;
                padding:12px 16px;margin-bottom:16px;font-size:0.82em;color:#8b8fa8'>
       <strong style='color:#f39c12'>Demo note:</strong>
      This deployment uses <strong>Qwen2.5-0.5B</strong> (CPU-optimised) for free hosting.
      Full results in the notebook use <strong>Mistral-7B</strong>.
      Corpus: 3,000 Wikipedia articles — try
      <em>Algae, Anarchism, Aristotle, Astronomical unit, Alexander Graham Bell</em>.
    </div>
    """)

    stats_display = gr.HTML(value=build_stats_html())

    with gr.Row():
        with gr.Column(scale=1):
            question_box = gr.Textbox(
                placeholder="e.g. Who was Aristotle?",
                lines=3,
                label="Ask a Question",
            )
            num_chunks = gr.Slider(
                minimum=1, maximum=5, value=3, step=1,
                label="Chunks to retrieve (k)",
            )
            submit_btn = gr.Button("Ask TruthRAG →", elem_classes=["ask-btn"])

            gr.Examples(
                examples=[
                    ["Who invented the telephone?", 3],
                    ["What is Anarchism?", 3],
                    ["What is an Astronomical unit?", 3],
                    ["What is Algae?", 3],
                    ["Who was Aristotle?", 3],
                ],
                inputs=[question_box, num_chunks],
                label="Try an example",
            )

        with gr.Column(scale=2):
            answer_html = gr.HTML(
                value="<div style='color:#8b8fa8;text-align:center;padding:60px'>"
                      "Your answer will appear here</div>",
                label="Answer Analysis",
            )
            with gr.Accordion(" Retrieved Source Chunks", open=False):
                sources_box = gr.Markdown()

    submit_btn.click(
        fn=truthrag_demo,
        inputs=[question_box, num_chunks],
        outputs=[answer_html, sources_box, stats_display],
    )

if __name__ == "__main__":
    demo.launch()