Create Evaluators

Evaluators

@@ -0,0 +1,93 @@
+import os
+import pickle
+import numpy as np
+import faiss
+import torch
+from datasets import load_dataset
+import evaluate
+
+# Import RAG setup and retrieval logic from app.py
+from app import setup_rag, retrieve
+
+
+def retrieval_recall(dataset, passages, embedder, index, k=20, rerank_k=None, num_samples=100):
+    """
+    Compute raw Retrieval Recall@k on the first num_samples examples.
+    If rerank_k is set, also apply cross-encoder reranking.
+    """
+    hits = 0
+    for ex in dataset.select(range(num_samples)):
+        question = ex["question"]
+        gold_answers = ex["answers"]["text"]
+        # get top-k retrieved contexts
+        if rerank_k:
+            ctxs, _ = retrieve(question, passages, embedder, index, k=k, rerank_k=rerank_k)
+        else:
+            # skip reranking: use top-k directly
+            q_emb = embedder.encode([question], convert_to_numpy=True)
+            distances, idxs = index.search(q_emb, k)
+            ctxs = [passages[i] for i in idxs[0]]
+        # check if any gold span appears
+        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, index, k=20, rerank_k=None, num_samples=100):
+    """
+    Retrieval Recall@k evaluated only on answerable questions.
+    """
+    hits, total = 0, 0
+    for ex in dataset.select(range(num_samples)):
+        if not ex["answers"]["text"]:
+            continue
+        total += 1
+        question = ex["question"]
+        if rerank_k:
+            ctxs, _ = retrieve(question, passages, embedder, 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 ex["answers"]["text"]):
+            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):
+    """
+    End-to-end QA EM/F1 on answerable subset using the retrieve_and_answer logic.
+    """
+    squad_metric = evaluate.load("squad")
+    preds, refs = [], []
+    for ex in dataset.select(range(num_samples)):
+        if not ex["answers"]["text"]:
+            continue
+        qid = ex["id"]
+        # retrieve and generate
+        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
+
+
+def main():
+    # Setup RAG components
+    passages, embedder, reranker, index, qa_pipe = setup_rag()
+    # Load SQuAD v2 validation set
+    squad = load_dataset("rajpurkar/squad_v2", split="validation")
+
+    # Run evaluations
+    retrieval_recall(squad, passages, embedder, index, k=20, rerank_k=5, num_samples=100)
+    retrieval_recall_answerable(squad, passages, embedder, index, k=20, rerank_k=5, num_samples=100)
+    qa_eval_answerable(squad, passages, embedder, reranker, index, qa_pipe, k=20, num_samples=100)
+
+
+if __name__ == "__main__":
+    main()