experiments/evaluate_beir.py

import sys
import os
import numpy as np
from datasets import load_dataset

# Add project root to python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.hybrid_search import BM25Index, HybridSearcher
from app.vector_store import build_index, search_index
from app.evaluation import precision_at_k, mrr, ndcg_at_k, recall_at_k, average_precision
from sentence_transformers import SentenceTransformer

try:
    from app.reranker import get_reranker
except ImportError:
    get_reranker = None

def run_beir_evaluation():
    print("Loading BEIR SciFact Dataset...")
    
    # Load corpus
    corpus = load_dataset('mteb/scifact', 'corpus', split='corpus')
    docs = []
    doc_id_to_idx = {}
    for i, item in enumerate(corpus):
        doc_text = f"{item['title']} {item['text']}"
        docs.append(doc_text.strip())
        doc_id_to_idx[item['_id']] = i
        
    print(f"Loaded {len(docs)} documents.")
    
    # Load queries
    queries = load_dataset('mteb/scifact', 'queries', split='queries')
    query_id_to_text = {item['_id']: item['text'] for item in queries}
    
    # Load qrels (ground truth)
    qrels = load_dataset('mteb/scifact', 'default', split='test')
    
    # Group qrels by query_id
    from collections import defaultdict
    benchmark_queries = defaultdict(list)
    for item in qrels:
        q_id = item['query-id']
        doc_id = item['corpus-id']
        score = item['score']
        if score > 0 and doc_id in doc_id_to_idx:
            benchmark_queries[q_id].append(doc_id_to_idx[doc_id])
            
    # Format into list
    eval_queries = []
    for q_id, rel_docs in benchmark_queries.items():
        if rel_docs and q_id in query_id_to_text:
            eval_queries.append({
                "query": query_id_to_text[q_id],
                "relevant_docs": rel_docs
            })
            
    eval_queries = eval_queries[:50]  # Back to 50 queries now that we use the fast S-PubMedBert model
    print(f"Loaded {len(eval_queries)} queries with ground truth.")
    
    # Initialize Search Components
    print("\nLoading Embedding Model...")
    model = SentenceTransformer('all-MiniLM-L6-v2')
    
    print("Encoding Documents (Dense Index)...")
    embeddings = model.encode(docs, show_progress_bar=True, convert_to_numpy=True, normalize_embeddings=True)
    
    # Fake labels for building the index
    labels = [0] * len(docs)
    faiss_index = build_index(embeddings, labels)
    
    print("\nBuilding BM25 Index (Sparse Index)...")
    bm25 = BM25Index()
    bm25.fit(docs)
    
    hybrid = HybridSearcher(bm25)
    
    reranker = None
    if get_reranker is not None:
        try:
            reranker = get_reranker()
            _ = reranker.model  # Force load
        except Exception as e:
            print(f"Failed to load reranker: {e}")
            
    print(f"\nEvaluating Search Modes using {len(eval_queries)} Natural Language Queries...")
    
    metrics = {
        "Dense": {"p@3": [], "mrr": [], "ndcg@10": [], "recall@10": [], "map": []},
        "BM25": {"p@3": [], "mrr": [], "ndcg@10": [], "recall@10": [], "map": []},
        "Hybrid": {"p@3": [], "mrr": [], "ndcg@10": [], "recall@10": [], "map": []},
        "Reranked": {"p@3": [], "mrr": [], "ndcg@10": [], "recall@10": [], "map": []}
    }
    
    for i, item in enumerate(eval_queries):
        print(f"Evaluating query {i+1}/{len(eval_queries)}...", end="\r")
        q_text = item["query"]
        relevant_docs = item["relevant_docs"]
        
        q_emb = model.encode([q_text], convert_to_numpy=True, normalize_embeddings=True)[0]
        
        # Dense Search
        dense_distances, dense_indices = search_index(faiss_index, q_emb, top_k=10)
        dense_indices = [int(idx) for idx in dense_indices if idx >= 0]
        if dense_indices:
            metrics["Dense"]["p@3"].append(precision_at_k(relevant_docs, dense_indices, 3))
            metrics["Dense"]["mrr"].append(mrr(relevant_docs, dense_indices))
            metrics["Dense"]["ndcg@10"].append(ndcg_at_k(relevant_docs, dense_indices, 10))
            metrics["Dense"]["recall@10"].append(recall_at_k(relevant_docs, dense_indices, 10))
            metrics["Dense"]["map"].append(average_precision(relevant_docs, dense_indices))
        
        # BM25 Search
        _, bm25_indices = bm25.score(q_text, top_k=10)
        bm25_indices = [int(idx) for idx in bm25_indices if idx >= 0]
        if bm25_indices:
            metrics["BM25"]["p@3"].append(precision_at_k(relevant_docs, bm25_indices, 3))
            metrics["BM25"]["mrr"].append(mrr(relevant_docs, bm25_indices))
            metrics["BM25"]["ndcg@10"].append(ndcg_at_k(relevant_docs, bm25_indices, 10))
            metrics["BM25"]["recall@10"].append(recall_at_k(relevant_docs, bm25_indices, 10))
            metrics["BM25"]["map"].append(average_precision(relevant_docs, bm25_indices))
        
        # Hybrid Search
        hybrid_indices, _, _ = hybrid.search(q_text, q_emb, faiss_index, docs, top_k=10)
        hybrid_indices = [int(idx) for idx in hybrid_indices if idx >= 0]
        if hybrid_indices:
            metrics["Hybrid"]["p@3"].append(precision_at_k(relevant_docs, hybrid_indices, 3))
            metrics["Hybrid"]["mrr"].append(mrr(relevant_docs, hybrid_indices))
            metrics["Hybrid"]["ndcg@10"].append(ndcg_at_k(relevant_docs, hybrid_indices, 10))
            metrics["Hybrid"]["recall@10"].append(recall_at_k(relevant_docs, hybrid_indices, 10))
            metrics["Hybrid"]["map"].append(average_precision(relevant_docs, hybrid_indices))
            
        # Reranked Hybrid (Top-20 retrieved -> Top-10 reranked with Interpolation)
        if reranker is not None:
            cand_indices, cand_scores, _ = hybrid.search(q_text, q_emb, faiss_index, docs, top_k=20)
            cand_indices = [int(idx) for idx in cand_indices if idx >= 0]
            if cand_indices:
                cand_docs = [docs[idx] for idx in cand_indices]
                r_indices, _ = reranker.rerank(q_text, cand_docs, cand_indices, hybrid_scores=cand_scores, top_k=10)
                metrics["Reranked"]["p@3"].append(precision_at_k(relevant_docs, r_indices, 3))
                metrics["Reranked"]["mrr"].append(mrr(relevant_docs, r_indices))
                metrics["Reranked"]["ndcg@10"].append(ndcg_at_k(relevant_docs, r_indices, 10))
                metrics["Reranked"]["recall@10"].append(recall_at_k(relevant_docs, r_indices, 10))
                metrics["Reranked"]["map"].append(average_precision(relevant_docs, r_indices))

    print("\n" + "="*70)
    print(f"{'Mode':<12} | {'P@3':<8} | {'MRR':<8} | {'NDCG@10':<8} | {'R@10':<8} | {'MAP':<8}")
    print("-" * 70)
    
    modes = ["BM25", "Dense", "Hybrid"]
    if reranker is not None:
        modes.append("Reranked")
        
    for mode in modes:
        if not metrics[mode]["p@3"]:
            continue
        p3 = np.mean(metrics[mode]["p@3"])
        mrr_val = np.mean(metrics[mode]["mrr"])
        ndcg = np.mean(metrics[mode]["ndcg@10"])
        rec = np.mean(metrics[mode]["recall@10"])
        m_ap = np.mean(metrics[mode]["map"])
        
        print(f"{mode:<12} | {p3:.4f}   | {mrr_val:.4f}   | {ndcg:.4f}   | {rec:.4f}   | {m_ap:.4f}")
    print("="*70)

if __name__ == '__main__':
    run_beir_evaluation()