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"""
{d['sentence']}
{d['sentence']}
{d['sentence']}
Retrieval-Augmented Generation with Automated Hallucination Detection