Update app.py

app.py

@@ -3,6 +3,7 @@
 
 import os
 import pickle
+import argparse
 import faiss
 import numpy as np
 import torch
@@ -15,6 +16,7 @@ from transformers import (
     AutoModelForSeq2SeqLM,
     pipeline as hf_pipeline,
 )
+import evaluate
 
 # ── 1. Configuration ──
 DATA_DIR       = os.path.join(os.getcwd(), "data")
@@ -30,7 +32,6 @@ 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"MODEL={MODEL_NAME}, EMBEDDER={EMBEDDER_MODEL}, DEVICE={'GPU' if DEVICE==0 else 'CPU'}")
 
 # ── 2. Helpers ──
 def make_context_snippets(contexts, max_words=MAX_CTX_WORDS):
@@ -53,15 +54,15 @@ def chunk_text(text, max_tokens, stride=None):
         start += stride
     return chunks
 
+
 # ── 3. Load & preprocess passages ──
 def load_passages():
-    # 3.1 load raw corpora
     wiki_ds   = load_dataset("rag-datasets/rag-mini-wikipedia", "text-corpus", split="passages")
     squad_ds  = load_dataset("rajpurkar/squad_v2", split="train[:100]")
     trivia_ds = load_dataset("mandarjoshi/trivia_qa", "rc", split="validation[:100]")
 
-    wiki_passages   = wiki_ds["passage"]
-    squad_passages  = [ex["context"] for ex in squad_ds]
+    wiki_passages  = wiki_ds["passage"]
+    squad_passages = [ex["context"] for ex in squad_ds]
     trivia_passages = []
     for ex in trivia_ds:
         for fld in ("wiki_context", "search_context"):
@@ -69,12 +70,10 @@ def load_passages():
             if txt:
                 trivia_passages.append(txt)
 
-    # dedupe
     all_passages = list(dict.fromkeys(wiki_passages + squad_passages + trivia_passages))
+    tokenizer   = AutoTokenizer.from_pretrained(MODEL_NAME)
+    max_tokens  = tokenizer.model_max_length
 
-    # 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)
@@ -88,6 +87,7 @@ def load_passages():
         pickle.dump(chunks, f)
     return chunks
 
+
 # ── 4. Build or load FAISS ──
 def load_faiss_index(passages):
     embedder = SentenceTransformer(EMBEDDER_MODEL)
@@ -116,6 +116,7 @@ def load_faiss_index(passages):
 
     return embedder, reranker, index
 
+
 # ── 5. Initialize RAG components ──
 def setup_rag():
     if os.path.exists(PCTX_PATH):
@@ -141,8 +142,9 @@ def setup_rag():
 
     return passages, embedder, reranker, index, qa_pipe
 
+
 # ── 6. Retrieval & generation ──
-def retrieve(question, passages, embedder, index, k=20, rerank_k=5):
+def retrieve(question, passages, embedder, reranker, index, k=20, rerank_k=5):
     q_emb      = embedder.encode([question], convert_to_numpy=True)
     distances, idxs = index.search(q_emb, k)
 
@@ -166,7 +168,7 @@ def generate(question, contexts, qa_pipe):
     return qa_pipe(prompt)[0]["generated_text"].strip()
 
 def retrieve_and_answer(question, passages, embedder, reranker, index, qa_pipe):
-    contexts, dists = retrieve(question, passages, embedder, index)
+    contexts, dists = retrieve(question, passages, embedder, reranker, index)
     if not contexts or dists[0] > DIST_THRESHOLD:
         return "Sorry, I don't know.", []
     return generate(question, contexts, qa_pipe), contexts
@@ -181,24 +183,99 @@ def answer_and_contexts(question, passages, embedder, reranker, index, qa_pipe):
     ]
     return ans, "\n\n---\n\n".join(snippets)
 
-# ── 7. Gradio app ──
+
+# ── 7. Evaluation routines ──
+def retrieval_recall(dataset, passages, embedder, reranker, index, k=20, rerank_k=None, num_samples=100):
+    hits = 0
+    for ex in dataset.select(range(num_samples)):
+        question     = ex["question"]
+        gold_answers = ex["answers"]["text"]
+
+        if rerank_k:
+            ctxs, _ = retrieve(question, passages, embedder, reranker, index, k=k, rerank_k=rerank_k)
+        else:
+            q_emb      = embedder.encode([question], convert_to_numpy=True)
+            distances, idxs = index.search(q_emb, k)
+            ctxs = [passages[i] for i in idxs[0]]
+
+        if any(any(ans in ctx for ctx in ctxs) for ans in gold_answers):
+            hits += 1
+
+    recall = hits / num_samples
+    print(f"Retrieval Recall@{k} (rerank_k={rerank_k}): {recall:.3f} ({hits}/{num_samples})")
+    return recall
+
+def retrieval_recall_answerable(dataset, passages, embedder, reranker, index, k=20, rerank_k=None, num_samples=100):
+    hits, total = 0, 0
+    for ex in dataset.select(range(num_samples)):
+        gold = ex["answers"]["text"]
+        if not gold:
+            continue
+        total += 1
+        question = ex["question"]
+
+        if rerank_k:
+            ctxs, _ = retrieve(question, passages, embedder, reranker, index, k=k, rerank_k=rerank_k)
+        else:
+            q_emb      = embedder.encode([question], convert_to_numpy=True)
+            distances, idxs = index.search(q_emb, k)
+            ctxs = [passages[i] for i in idxs[0]]
+
+        if any(any(ans in ctx for ctx in ctxs) for ans in gold):
+            hits += 1
+
+    recall = hits / total if total > 0 else 0.0
+    print(f"Retrieval Recall@{k} on answerable only (rerank_k={rerank_k}): {recall:.3f} ({hits}/{total})")
+    return recall
+
+def qa_eval_answerable(dataset, passages, embedder, reranker, index, qa_pipe, k=20, num_samples=100):
+    squad_metric = evaluate.load("squad")
+    preds, refs  = [], []
+    for ex in dataset.select(range(num_samples)):
+        gold = ex["answers"]["text"]
+        if not gold:
+            continue
+        qid = ex["id"]
+        answer, _ = retrieve_and_answer(ex["question"], passages, embedder, reranker, index, qa_pipe)
+        preds.append({"id": qid, "prediction_text": answer})
+        refs.append({"id": qid, "answers": ex["answers"]})
+
+    results = squad_metric.compute(predictions=preds, references=refs)
+    print(f"Answerable-only QA EM: {results['exact_match']:.2f}, F1: {results['f1']:.2f}")
+    return results
+
+
+# ── 8. Main entry ──
 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'?"
-        ],
-        allow_flagging="never",
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--eval", action="store_true",
+        help="Run retrieval/QA evaluations on SQuAD instead of launching the demo"
     )
-    demo.launch()
+    args = parser.parse_args()
+
+    if args.eval:
+        squad = load_dataset("rajpurkar/squad_v2", split="validation")
+        retrieval_recall(squad, passages, embedder, reranker, index, k=20, rerank_k=5, num_samples=100)
+        retrieval_recall_answerable(squad, passages, embedder, reranker, index, k=20, rerank_k=5, num_samples=100)
+        qa_eval_answerable(squad, passages, embedder, reranker, index, qa_pipe, k=20, num_samples=100)
+    else:
+        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'?"
+            ],
+            allow_flagging="never",
+        )
+        demo.launch()
 
 if __name__ == "__main__":
     main()