beir.py

import psycopg2
import psycopg2.extras
import time
import numpy as np
from collections import defaultdict
from datasets import load_dataset

DB_CONFIG = {
    'dbname': 'bible_dev',
    'user': 'postgres',
    'password': 'postgres',
    'host': 'localhost',
    'port': 5432
}

DATASET_NAME = 'nfcorpus'
TOP_K = 10


def connect_db():
    return psycopg2.connect(**DB_CONFIG)


def load_beir_dataset(dataset_name):
    print(f"Loading BEIR dataset: {dataset_name}")

    corpus = load_dataset(f'BeIR/{dataset_name}', 'corpus')
    docs = {row['_id']: {'title': row['title'], 'text': row['text']}
            for row in corpus['corpus']}

    queries_dataset = load_dataset(f'BeIR/{dataset_name}', 'queries')
    queries = {row['_id']: row['text'] for row in queries_dataset['queries']}

    # relevance judgments dataset
    qrels_dataset = load_dataset(f'BeIR/{dataset_name}-qrels')
    qrels = defaultdict(dict)
    for row in qrels_dataset['validation']:
        qrels[row['query-id']][row['corpus-id']] = row['score']

    return docs, queries, qrels


def index_documents(conn, docs):
    """Index all documents into PostgreSQL"""
    print(f"Indexing {len(docs)} documents...")
    cur = conn.cursor()

    cur.execute("TRUNCATE verses, verse_stats, term_stats CASCADE")

    doc_id_map = {}
    batch = []

    for idx, (doc_id, content) in enumerate(docs.items(), start=1):
        full_text = f"{content['title']} {content['text']}"
        numeric_id = idx
        doc_id_map[doc_id] = numeric_id
        batch.append((numeric_id, full_text))

        if len(batch) >= 1000:
            psycopg2.extras.execute_batch(
                cur,
                "INSERT INTO verses (id, text, inserted_at, updated_at) VALUES (%s, %s, now(), now())",
                batch
            )
            batch = []

    if batch:
        psycopg2.extras.execute_batch(
            cur,
            "INSERT INTO verses (id, text, inserted_at, updated_at) VALUES (%s, %s, now(), now())",
            batch
        )

    conn.commit()
    print("Documents inserted, now indexing for BM25...")

    cur.execute("""
        SELECT FROM index_all_verses()
    """)

    conn.commit()
    cur.close()

    print("Indexing complete!")

    return doc_id_map


def search_bm25(conn, query, k=10):
    cur = conn.cursor()
    cur.execute("""
        SELECT verse_id, score, content
        FROM search_verses(%s, 1.2, 0.75, %s, 0.3)
    """, (query, k))

    results = cur.fetchall()
    cur.close()
    return results


def search_ilike(conn, query, k=10):
    cur = conn.cursor()

    pattern = f"%{query}%"
    cur.execute("""
        SELECT id, 0 as score, text
        FROM verses
        WHERE text ILIKE %s
        LIMIT %s
    """, (pattern, k))

    results = cur.fetchall()
    cur.close()
    return results


def calculate_dcg(relevances, k):
    """Calculate Discounted Cumulative Gain"""
    relevances = np.array(relevances[:k])
    if len(relevances) == 0:
        return 0.0

    discounts = np.log2(np.arange(len(relevances)) + 2)
    return np.sum(relevances / discounts)


def calculate_ndcg(retrieved_docs, relevant_docs, k=10):
    """Calculate Normalized Discounted Cumulative Gain @ k"""
    relevances = []
    for doc_id in retrieved_docs[:k]:
        relevances.append(relevant_docs.get(str(doc_id), 0))

    dcg = calculate_dcg(relevances, k)

    ideal_relevances = sorted(relevant_docs.values(), reverse=True)
    idcg = calculate_dcg(ideal_relevances, k)

    if idcg == 0:
        return 0.0

    return dcg / idcg


def calculate_precision_at_k(retrieved_doc_ids, relevant_doc_ids_for_query, k):
    """
    Calculate Precision @ k.
    Proportion of retrieved documents that are relevant.
    """
    if not retrieved_doc_ids:
        return 0.0

    retrieved_k = retrieved_doc_ids[:k]
    num_relevant_retrieved = sum(1 for doc_id in retrieved_k if str(doc_id) in relevant_doc_ids_for_query)
    return num_relevant_retrieved / len(retrieved_k)


def calculate_recall_at_k(retrieved_doc_ids, relevant_doc_ids_for_query, k):
    """
    Calculate Recall @ k.
    Proportion of relevant documents found within the top k retrieved.
    """
    num_relevant_docs_in_qrels = len(relevant_doc_ids_for_query)
    if num_relevant_docs_in_qrels == 0:
        return 0.0

    retrieved_k = retrieved_doc_ids[:k]
    num_relevant_retrieved = sum(1 for doc_id in retrieved_k if str(doc_id) in relevant_doc_ids_for_query)
    return num_relevant_retrieved / num_relevant_docs_in_qrels


def run_benchmark(conn, queries, qrels, search_func, method_name, doc_id_map, k=10):
    """Run benchmark for a single search method"""
    print(f"\nRunning benchmark for {method_name}...")

    precision_scores = []
    recall_scores = []
    ndcg_scores = []
    total_time = 0
    query_count = 0
    wins = 0  # queries that returned at least one relevant result

    # Create reverse mapping (numeric_id -> beir_id) for results
    reverse_map = {v: k for k, v in doc_id_map.items()}

    for query_id, query_text in queries.items():
        if query_id not in qrels:
            continue

        relevant_doc_ids_for_query = {doc_id for doc_id, score in qrels[query_id].items() if score > 0}
        if not relevant_doc_ids_for_query:
            continue

        start_time = time.time()
        results = search_func(conn, query_text, k)
        query_time = time.time() - start_time

        total_time += query_time
        query_count += 1

        retrieved_docs = [reverse_map.get(row[0], str(row[0])) for row in results]

        ndcg = calculate_ndcg(retrieved_docs, qrels[query_id], k)
        ndcg_scores.append(ndcg)

        precision = calculate_precision_at_k(retrieved_docs, relevant_doc_ids_for_query, k)
        precision_scores.append(precision)

        recall = calculate_recall_at_k(retrieved_docs, relevant_doc_ids_for_query, k)
        recall_scores.append(recall)

        # Count as a "win" if we found any relevant documents
        if ndcg > 0:
            wins += 1

        if query_count % 100 == 0:
            print(f"  Processed {query_count} queries...")

    avg_ndcg = np.mean(ndcg_scores) if ndcg_scores else 0
    avg_precision = np.mean(precision_scores) if precision_scores else 0
    avg_recall = np.mean(recall_scores) if recall_scores else 0
    qps = query_count / total_time if total_time > 0 else 0
    avg_latency_ms = (total_time / query_count * 1000) if query_count > 0 else 0
    recall = (wins / query_count * 100) if query_count > 0 else 0

    return {
        'method': method_name,
        'precision@10': avg_precision,
        'recall@10': avg_recall,
        'ndcg@10': avg_ndcg,
        'qps': qps,
        'avg_latency_ms': avg_latency_ms,
        'total_queries': query_count,
        'total_time': total_time,
        'recall_pct': recall
    }


def print_comparison(bm25_results, ilike_results):
    """Print side-by-side comparison table"""
    print("\n" + "=" * 80)
    print("BENCHMARK RESULTS COMPARISON")
    print("=" * 80)
    print(f"Dataset: {DATASET_NAME} | Total Queries: {bm25_results['total_queries']}")
    print("-" * 80)
    print(f"{'Metric':<25} | {'BM25':<15} | {'ILIKE':<15} | {'Improvement':<15}")
    print("-" * 80)

    # NDCG@10
    ndcg_improvement = ((bm25_results['ndcg@10'] - ilike_results['ndcg@10']) /
                        ilike_results['ndcg@10'] * 100) if ilike_results['ndcg@10'] > 0 else float('inf')
    print(f"{'NDCG@10 (relevance)':<25} | {bm25_results['ndcg@10']:<15.4f} | "
          f"{ilike_results['ndcg@10']:<15.4f} | {ndcg_improvement:>+14.1f}%")

    # QPS
    qps_improvement = ((bm25_results['qps'] - ilike_results['qps']) /
                       ilike_results['qps'] * 100) if ilike_results['qps'] > 0 else 0
    print(f"{'QPS (throughput)':<25} | {bm25_results['qps']:<15.2f} | "
          f"{ilike_results['qps']:<15.2f} | {qps_improvement:>+14.1f}%")

    # Latency (lower is better, so invert the improvement)
    latency_improvement = ((ilike_results['avg_latency_ms'] - bm25_results['avg_latency_ms']) /
                           ilike_results['avg_latency_ms'] * 100) if ilike_results['avg_latency_ms'] > 0 else 0
    print(f"{'Avg Latency (ms)':<25} | {bm25_results['avg_latency_ms']:<15.2f} | "
          f"{ilike_results['avg_latency_ms']:<15.2f} | {latency_improvement:>+14.1f}%")

    # Precision@10
    precision_improvement = ((bm25_results['precision@10'] - ilike_results['precision@10']) /
                             ilike_results['precision@10'] * 100) if ilike_results['precision@10'] > 0 else float('inf')
    print(f"{'Precision@10':<25} | {bm25_results['precision@10']:<15.4f} | "
          f"{ilike_results['precision@10']:<15.4f} | {precision_improvement:>+14.1f}%")

    # Recall@10
    # Note: For recall, if ilike_results['recall@10'] is 0, improvement can be tricky to display as %
    # We'll just display the difference in percentage points (pp)
    recall_abs_diff = bm25_results['recall@10'] - ilike_results['recall@10']
    print(f"{'Recall@10':<25} | {bm25_results['recall@10']:<15.4f} | "
          f"{ilike_results['recall@10']:<15.4f} | {recall_abs_diff:>+14.4f}pp")

    # Recall percentage
    recall_improvement = bm25_results['recall_pct'] - ilike_results['recall_pct']
    print(f"{'Recall % (found relevant)':<25} | {bm25_results['recall_pct']:<15.1f} | "
          f"{ilike_results['recall_pct']:<15.1f} | {recall_improvement:>+14.1f}pp")

    print("-" * 80)
    print(f"{'Total Time':<25} | {bm25_results['total_time']:<15.2f} | "
          f"{ilike_results['total_time']:<15.2f} | (seconds)")
    print("=" * 80)

    print("\nKEY TAKEAWAYS:")
    if ndcg_improvement > 0:
        print(f"  ✓ BM25 provides {ndcg_improvement:.1f}% better relevance ranking")
    if bm25_results['precision@10'] > ilike_results['precision@10']:
        print(f"  ✓ BM25 retrieves more relevant documents within the top {TOP_K} (Precision@10)")
    if bm25_results['recall@10'] > ilike_results['recall@10']:
        print(f"  ✓ BM25 finds a larger proportion of all relevant documents (Recall@10)")
    if qps_improvement > 0:
        print(f"  ✓ BM25 is {qps_improvement:.1f}% faster (higher throughput)")
    elif qps_improvement < 0:
        print(f"  ⚠ BM25 is {abs(qps_improvement):.1f}% slower, but much more relevant")
    if recall_improvement > 0:
        print(f"  ✓ BM25 finds relevant results for {recall_improvement:.1f}% more queries")
    print()


def main():
    """Main benchmark execution"""
    print("=" * 80)
    print("BEIR Benchmark: BM25 vs ILIKE Comparison")
    print("=" * 80)

    docs, queries, qrels = load_beir_dataset(DATASET_NAME)
    print(f"Loaded: {len(docs)} documents, {len(queries)} queries, {len(qrels)} query pairs")

    conn = connect_db()

    try:
        doc_id_map = index_documents(conn, docs)

        bm25_results = run_benchmark(conn, queries, qrels, search_bm25, "BM25", doc_id_map, k=TOP_K)
        ilike_results = run_benchmark(conn, queries, qrels, search_ilike, "ILIKE", doc_id_map, k=TOP_K)

        print_comparison(bm25_results, ilike_results)

    finally:
        conn.close()


if __name__ == "__main__":
    main()