inspect_rag_hf.py

"""inspect_rag_hf.py — Inspect HF RAG results on random test cases (no GPU / no server)."""
import json
import random
import logging
import argparse
from pathlib import Path
from rag_hf import FAISSRetrieverHF

logging.basicConfig(level=logging.ERROR)
logging.getLogger("rag_hf").setLevel(logging.INFO)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n",    type=int, default=5,    help="Number of cases to inspect")
    parser.add_argument("--seed", type=int, default=None, help="Random seed")
    parser.add_argument("--topk", type=int, default=5,    help="Top-K to retrieve")
    args = parser.parse_args()

    retriever = FAISSRetrieverHF()
    print("Loading HF RAG index (from cache)…")
    retriever.load()
    print("RAG ready.\n")

    test_dir = Path("data/test_set")
    test_files = list(test_dir.glob("*.json"))
    sample_files = random.Random(args.seed).sample(test_files, min(args.n, len(test_files)))

    print("=" * 100)
    print(f"HF RAG INSPECTION: {len(sample_files)} RANDOM TEST CASES (top-{args.topk})")
    print("=" * 100)

    for i, f in enumerate(sample_files, 1):
        data      = json.loads(f.read_text())
        query     = data["query"]
        gt        = data["gt"]
        valid     = set(data["icd_codes"])

        print(f"\n[{i}/{len(sample_files)}] FILE: {f.name}")
        print(f"SYMPTOMS  : {query[:400]}")
        print(f"GT CODE   : {gt}   (valid set: {sorted(valid)})")
        print("-" * 60)

        results = retriever.retrieve(query, top_k=args.topk)
        print(f"TOP {args.topk} RAG RESULTS:")
        for j, r in enumerate(results, 1):
            hit = "✅" if any(c in valid for c in r.icd_codes) else "❌"
            print(f"  {j}. {hit} [{', '.join(r.icd_codes)}]  {r.title[:80]}")
            print(f"     score={r.score:.4f}")
        print("=" * 100)


if __name__ == "__main__":
    main()