Chapter 9: Evaluation

# A labeled test set: (question, correct section) pairs, roughly two
# per section, including trap questions from Chapters 2, 5, and 7.
test_set = [
    ("how many vacation days can be carried over into next year?", "Vacation Policy"),
    ("do unused holidays roll over to next year?", "Vacation Policy"),
    ("how many sick days can be carried over to next year?", "Sick Leave Policy"),
    ("do i need a doctor's note if i am sick for several days?", "Sick Leave Policy"),
    ("how long do i have to submit an expense report?", "Expense Policy"),
    ("who reviews my purchase receipts?", "Expense Policy"),
    ("how much is the home office stipend?", "Remote Work Policy"),
    ("what equipment does the stipend cover?", "Remote Work Policy"),
    ("what are the standard working hours?", "Working Hours"),
    ("do i need approval before working overtime?", "Working Hours"),
    ("how should i treat my colleagues at work?", "Code of Conduct"),
    ("what happens if i violate the code of conduct?", "Code of Conduct"),
    ("what department is 4477", "IT Support"),
    ("who can help me reset my password?", "IT Support"),
]
print(f"{len(test_set)} labeled test questions, covering all {len(section_chunks)} sections.")

def top_k(question, score_fn, k):
    return sorted(section_chunks, key=lambda c: score_fn(question, c), reverse=True)[:k]

def recall_at_k(score_fn, k=1):
    hits = 0
    for question, expected in test_set:
        titles = [section_title(c) for c in top_k(question, score_fn, k)]
        hits += expected in titles
    return hits / len(test_set)

def word_overlap_score(question, candidate):
    q_words = set(question.lower().split())
    c_words = set(candidate.lower().split())
    return len(q_words & c_words)

def bi_encoder_score(question, candidate):
    return cosine_similarity(embed_text(question), embed_text(candidate))

print(f"word overlap (Ch1):  recall@1 = {recall_at_k(word_overlap_score):.3f}")
print(f"embeddings (Ch2-6):  recall@1 = {recall_at_k(bi_encoder_score):.3f}")

class BM25:
    def __init__(self, documents, k1=1.5, b=0.75):
        self.k1, self.b = k1, b
        self.doc_tokens = [re.findall(r"[a-z0-9\']+", d.lower()) for d in documents]
        self.doc_lens = [len(t) for t in self.doc_tokens]
        self.avgdl = sum(self.doc_lens) / len(self.doc_lens)
        self.N = len(documents)
        doc_freq = {}
        for toks in self.doc_tokens:
            for term in set(toks):
                doc_freq[term] = doc_freq.get(term, 0) + 1
        self.idf = {t: math.log((self.N - n + 0.5) / (n + 0.5) + 1) for t, n in doc_freq.items()}

    def score(self, query, doc_index):
        term_freq = {}
        for t in self.doc_tokens[doc_index]:
            term_freq[t] = term_freq.get(t, 0) + 1
        dl = self.doc_lens[doc_index]
        total = 0.0
        for term in re.findall(r"[a-z0-9']+", query.lower()):
            if term not in self.idf:
                continue
            f = term_freq.get(term, 0)
            total += self.idf[term] * (f * (self.k1 + 1)) / (f + self.k1 * (1 - self.b + self.b * dl / self.avgdl))
        return total

bm25 = BM25(section_chunks)

def rrf_top_k(question, k, rrf_k=60):
    embed_ranked = sorted(range(len(section_chunks)), key=lambda i: bi_encoder_score(question, section_chunks[i]), reverse=True)
    bm25_ranked = sorted(range(len(section_chunks)), key=lambda i: bm25.score(question, i), reverse=True)
    scores = {}
    for rank, i in enumerate(embed_ranked, start=1):
        scores[i] = scores.get(i, 0) + 1 / (rrf_k + rank)
    for rank, i in enumerate(bm25_ranked, start=1):
        scores[i] = scores.get(i, 0) + 1 / (rrf_k + rank)
    ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)
    return [section_chunks[i] for i, _ in ranked[:k]]

def late_interaction_score(question, candidate):
    q_words = [w for w in tokenize(question) if w in word_index]
    c_words = [w for w in tokenize(candidate) if w in word_index]
    if not q_words or not c_words:
        return 0.0
    candidate_vectors = [embedding(w) for w in c_words]
    return sum(max(cosine_similarity(embedding(qw), cv) for cv in candidate_vectors) for qw in q_words) / len(q_words)

def hybrid_rerank_top_k(question, k, shortlist_size=3):
    shortlist = rrf_top_k(question, shortlist_size)
    reranked = sorted(shortlist, key=lambda c: late_interaction_score(question, c), reverse=True)
    return reranked[:k]

def recall_at_k_custom(topk_fn, k=1):
    hits = 0
    for question, expected in test_set:
        titles = [section_title(c) for c in topk_fn(question, k)]
        hits += expected in titles
    return hits / len(test_set)

print(f"word overlap (Ch1):      recall@1 = {recall_at_k(word_overlap_score):.3f}")
print(f"embeddings (Ch2-6):      recall@1 = {recall_at_k(bi_encoder_score):.3f}")
print(f"hybrid (Ch7):            recall@1 = {recall_at_k_custom(rrf_top_k):.3f}")
print(f"hybrid + rerank (Ch7+8): recall@1 = {recall_at_k_custom(hybrid_rerank_top_k):.3f}")

# Find a question hybrid alone got right, that hybrid+rerank got wrong.
for question, expected in test_set:
    hybrid_got = section_title(rrf_top_k(question, 1)[0])
    reranked_got = section_title(hybrid_rerank_top_k(question, 1)[0])
    if hybrid_got == expected and reranked_got != expected:
        print(f"Question: {question!r}")
        print(f"  hybrid alone got it right: {hybrid_got!r}")
        print(f"  hybrid + rerank got it wrong: {reranked_got!r}")
        print()
        shortlist = rrf_top_k(question, 3)
        print("  Re-ranker's view of the shortlist:")
        for c in shortlist:
            print(f"    {late_interaction_score(question, c):.3f}  {section_title(c)}")
        break

print(f"word overlap (Ch1):      recall@3 = {recall_at_k(word_overlap_score, k=3):.3f}")
print(f"embeddings (Ch2-6):      recall@3 = {recall_at_k(bi_encoder_score, k=3):.3f}")
print(f"hybrid (Ch7):            recall@3 = {recall_at_k_custom(rrf_top_k, k=3):.3f}")

print(f"word overlap:      {recall_at_k(word_overlap_score):.3f}")
print(f"embeddings:        {recall_at_k(bi_encoder_score):.3f}")
print(f"hybrid:            {recall_at_k_custom(rrf_top_k):.3f}")
print(f"hybrid + rerank:   {recall_at_k_custom(hybrid_rerank_top_k):.3f}")