app.py

#!/usr/bin/env python
# coding: utf-8

import os
import pickle
import faiss
import numpy as np
import torch
import gradio as gr

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, CrossEncoder
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    pipeline as hf_pipeline,
)

# ── 1. Configuration ──
DATA_DIR       = os.path.join(os.getcwd(), "data")
INDEX_PATH     = os.path.join(DATA_DIR, "faiss_index.faiss")
EMB_PATH       = os.path.join(DATA_DIR, "embeddings.npy")
PCTX_PATH      = os.path.join(DATA_DIR, "passages.pkl")

MODEL_NAME     = os.getenv("MODEL_NAME", "google/flan-t5-small")
EMBEDDER_MODEL = os.getenv("EMBEDDER_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
DIST_THRESHOLD = float(os.getenv("DIST_THRESHOLD", 1.0))
MAX_CTX_WORDS  = int(os.getenv("MAX_CTX_WORDS", 200))

DEVICE = 0 if torch.cuda.is_available() else -1

os.makedirs(DATA_DIR, exist_ok=True)

print(f"Using MODEL_NAME={MODEL_NAME}, EMBEDDER_MODEL={EMBEDDER_MODEL}, device={'GPU' if DEVICE==0 else 'CPU'}")

# ── 2. Helpers ──
def make_context_snippets(contexts, max_words=MAX_CTX_WORDS):
    out = []
    for c in contexts:
        words = c.split()
        if len(words) > max_words:
            c = " ".join(words[:max_words]) + " ... [truncated]"
        out.append(c)
    return out

def chunk_text(text, max_tokens, stride=None):
    words = text.split()
    if stride is None:
        stride = max_tokens // 4
    chunks, start = [], 0
    while start < len(words):
        end = start + max_tokens
        chunks.append(" ".join(words[start:end]))
        start += stride
    return chunks

# ── 3. Load & preprocess passages ──
def load_passages():
    # 3.1 load raw corpora
    wiki = load_dataset("rag-datasets/rag-mini-wikipedia", "text-corpus", split="passages")["passage"]
    squad = load_dataset("rajpurkar/squad_v2", split="train[:100]")["context"]
    trivia_ds = load_dataset("mandarjoshi/trivia_qa", "rc", split="validation[:100]")
    trivia = []
    for ex in trivia_ds:
        for fld in ("wiki_context", "search_context"):
            txt = ex.get(fld) or ""
            if txt: trivia.append(txt)

    all_passages = list(dict.fromkeys(wiki + squad + trivia))
    # 3.2 chunk long passages
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    max_tokens = tokenizer.model_max_length

    chunks = []
    for p in all_passages:
        toks = tokenizer.tokenize(p)
        if len(toks) > max_tokens:
            chunks.extend(chunk_text(p, max_tokens))
        else:
            chunks.append(p)

    print(f"[load_passages] total chunks: {len(chunks)}")
    with open(PCTX_PATH, "wb") as f:
        pickle.dump(chunks, f)
    return chunks

# ── 4. Build or load FAISS ──
def load_faiss_index(passages):
    # sentence‐transformers embedder + cross‐encoder
    embedder = SentenceTransformer(EMBEDDER_MODEL)
    reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

    if os.path.exists(INDEX_PATH) and os.path.exists(EMB_PATH):
        print("Loading FAISS index & embeddings from disk …")
        index = faiss.read_index(INDEX_PATH)
        embeddings = np.load(EMB_PATH)
    else:
        print("Encoding passages & building FAISS index …")
        embeddings = embedder.encode(passages, show_progress_bar=True, convert_to_numpy=True, batch_size=32)
        embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

        dim = embeddings.shape[1]
        index = faiss.IndexFlatIP(dim)
        index.add(embeddings)

        faiss.write_index(index, INDEX_PATH)
        np.save(EMB_PATH, embeddings)

    return embedder, reranker, index

# ── 5. Set up RAG pipeline ──
def setup_rag():
    # 5.1 load or build index + embedder/reranker
    if os.path.exists(PCTX_PATH):
        with open(PCTX_PATH, "rb") as f:
            passages = pickle.load(f)
    else:
        passages = load_passages()

    embedder, reranker, index = load_faiss_index(passages)

    # 5.2 load generator model & HF pipeline
    tok = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
    qa_pipe = hf_pipeline(
        "text2text-generation",
        model=model,
        tokenizer=tok,
        device=DEVICE,
        truncation=True,
        max_length=512,
        num_beams=4,        # optional: enable beam search
        early_stopping=True
    )

    return passages, embedder, reranker, index, qa_pipe

# ── 6. Retrieval + Generation ──
def retrieve(question, passages, embedder, index, k=20, rerank_k=5):
    q_emb = embedder.encode([question], convert_to_numpy=True)
    distances, idxs = index.search(q_emb, k)

    cands = [passages[i] for i in idxs[0]]
    scores = reranker.predict([[question, c] for c in cands])
    top = np.argsort(scores)[-rerank_k:][::-1]

    final_ctxs = [cands[i] for i in top]
    final_dists = [distances[0][i] for i in top]
    return final_ctxs, final_dists

def generate(question, contexts, qa_pipe):
    lines = [ f"Context {i+1}: {s}" 
              for i,s in enumerate(make_context_snippets(contexts)) ]
    prompt = (
        "You are a helpful assistant. Use ONLY the following contexts to answer. "
        "If the answer is not contained, say 'Sorry, I don't know.'\n\n"
        + "\n".join(lines)
        + f"\n\nQuestion: {question}\nAnswer:"
    )
    return qa_pipe(prompt)[0]["generated_text"].strip()

def retrieve_and_answer(question, passages, embedder, reranker, index, qa_pipe):
    ctxs, dists = retrieve(question, passages, embedder, index)
    if not ctxs or dists[0] > DIST_THRESHOLD:
        return "Sorry, I don't know.", []
    ans = generate(question, ctxs, qa_pipe)
    return ans, ctxs

def answer_and_contexts(question, 
                        passages, embedder, reranker, index, qa_pipe):
    ans, ctxs = retrieve_and_answer(question, passages, embedder, reranker, index, qa_pipe)
    if not ctxs:
        return ans, ""
    snippets = [
        f"Context {i+1}: {s}" 
        for i,s in enumerate(make_context_snippets(ctxs))
    ]
    return ans, "\n\n---\n\n".join(snippets)

# ── 7. Gradio app ──
def main():
    passages, embedder, reranker, index, qa_pipe = setup_rag()

    demo = gr.Interface(
        fn=lambda q: answer_and_contexts(q, passages, embedder, reranker, index, qa_pipe),
        inputs=gr.Textbox(lines=1, placeholder="Ask me anything…", label="Question"),
        outputs=[gr.Textbox(label="Answer"), gr.Textbox(label="Contexts")],
        title="🔍 RAG QA Demo",
        description="Retrieval-Augmented QA with threshold and context preview",
        examples=[
            "When was Abraham Lincoln inaugurated?",
            "What is the capital of France?",
            "Who wrote '1984'?"
        ]
    )
    demo.launch()

if __name__ == "__main__":
    main()