scripts/test_retrieval.py

"""
Day 2 retrieval quality test.

Compares dense-only vs BM25-only vs hybrid retrieval on 5 varied claims.
Prints a side-by-side comparison so you can see the improvement.

Run after the server has loaded (takes 1-2 min).
Usage:
    python scripts/test_retrieval.py
"""
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.pipeline import startup, retrieve_only, is_ready, _state


TEST_CLAIMS = [
    "Hypocretin neuron loss is associated with narcolepsy",
    "Aspirin reduces the risk of cardiovascular events",
    "Smoking increases lung cancer risk",
    "Vitamin D supplementation prevents COVID-19",
    "SARS-CoV-2 uses ACE2 receptor for cell entry",
]


def print_results(r, header):
    print(f"\n   {header}")
    print(f"   {'-' * 60}")
    for res in r["results"][:5]:
        title = (res['title'] or '(no title)')[:50]
        print(f"   [{res['rank']:2d}] score={res['score']:.3f} | doc={res['doc_id']:>6s} | {title}")


def main():
    if not is_ready():
        print("Loading pipeline (takes 1-2 min)...")
        startup()

    print("\n" + "=" * 70)
    print("  DAY 2 RETRIEVAL QUALITY TEST")
    print("  Comparing dense / BM25 / hybrid on 5 varied claims")
    print("=" * 70)

    for i, claim in enumerate(TEST_CLAIMS, 1):
        print(f"\n\n╔══ CLAIM {i}/{len(TEST_CLAIMS)}" + "═" * 52)
        print(f"║  {claim}")
        print("╚" + "═" * 68)

        dense = retrieve_only(claim, top_k=10, mode="dense")
        bm25 = retrieve_only(claim, top_k=10, mode="bm25")
        hybrid = retrieve_only(claim, top_k=10, mode="hybrid")

        print_results(dense,  "DENSE-ONLY  (pure vector search)")
        print_results(bm25,   "BM25-ONLY   (pure keyword search)")
        print_results(hybrid, "HYBRID      (dense + BM25 combined)")

        dense_docs  = set(r["doc_id"] for r in dense["results"][:5])
        bm25_docs   = set(r["doc_id"] for r in bm25["results"][:5])
        hybrid_docs = set(r["doc_id"] for r in hybrid["results"][:5])

        print(f"\n   Overlap top-5:")
        print(f"     dense ∩ bm25   = {len(dense_docs & bm25_docs)}/5 docs")
        print(f"     dense ∩ hybrid = {len(dense_docs & hybrid_docs)}/5 docs")
        print(f"     bm25  ∩ hybrid = {len(bm25_docs & hybrid_docs)}/5 docs")

    print("\n\n" + "=" * 70)
    print("  TEST COMPLETE")
    print("=" * 70)


if __name__ == "__main__":
    main()