Evaluating retrieval

import re
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer

sections = [p.strip() for p in re.split(r"\n(?=## )", PROTOCOL)[1:]]
ids = ["glucose", "kidney", "foot", "hypo", "eye"]

tests = [   # (question, id of the section that answers it)
    ("How often should HbA1c be checked after changing treatment?", "glucose"),
    ("What HbA1c target suits most adults?", "glucose"),
    ("My sugar dropped really low, what now?", "hypo"),
    ("What should I do about a hypo?", "hypo"),
    ("How many grams of carbohydrate treat a low?", "hypo"),
    ("Is the diabetes pill safe with weak kidneys?", "kidney"),
    ("When do we need to stop the biguanide?", "kidney"),
    ("How often is urine albumin checked?", "kidney"),
    ("Numb toes test", "foot"),
    ("When is a foot ulcer an emergency?", "foot"),
    ("How often are retinal photos needed?", "eye"),
    ("When should eye screening start?", "eye"),
]
questions = [q for q, _ in tests]
relevant = [ids.index(r) for _, r in tests]

vec = TfidfVectorizer(stop_words="english").fit(sections + questions)
keyword = cosine_similarity(vec.transform(questions), vec.transform(sections))
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
embedding = (model.encode(questions, normalize_embeddings=True)
             @ model.encode(sections, normalize_embeddings=True).T)

def ranks(scores):
    """1-based rank of the correct section for each question."""
    order = np.argsort(-scores, axis=1)
    return np.array([list(order[i]).index(relevant[i]) + 1 for i in range(len(relevant))])

def report(name, scores):
    r = ranks(scores)
    print(f"{name:<10} recall@1 {np.mean(r <= 1):.2f}   recall@3 {np.mean(r <= 3):.2f}   "
          f"MRR {np.mean(1 / r):.2f}")

def reciprocal_rank_fusion(*score_sets, k=60):
    """Hybrid search: add 1 / (k + rank) from each method."""
    fused = np.zeros_like(score_sets[0])
    for s in score_sets:
        order = np.argsort(-s, axis=1)
        for i in range(s.shape[0]):
            for rank, j in enumerate(order[i], start=1):
                fused[i, j] += 1 / (k + rank)
    return fused

report("Keyword", keyword)
report("Embedding", embedding)
report("Hybrid", reciprocal_rank_fusion(keyword, embedding))

# Keyword    recall@1 0.83   recall@3 0.92   MRR 0.88
# Embedding  recall@1 0.92   recall@3 1.00   MRR 0.94
# Hybrid     recall@1 0.92   recall@3 0.92   MRR 0.94

for name, s in [("Keyword", keyword), ("Embedding", embedding),
                ("Hybrid", reciprocal_rank_fusion(keyword, embedding))]:
    r = ranks(s)
    misses = [(questions[i], int(r[i])) for i in np.where(r > 1)[0]]
    print(name, "misses:", misses)

# Keyword misses: [('What should I do about a hypo?', 4), ('Is the diabetes pill safe with weak kidneys?', 3)]
# Embedding misses: [('What should I do about a hypo?', 3)]
# Hybrid misses: [('What should I do about a hypo?', 4)]