S118 · RAG II — Retrieval, Hybrid Search, Reranking

"""rag_retrieve.py — hybrid retrieval + cross-encoder reranking.
 
Assumes S117's Chroma collection 'docs_v1' already exists.
Requires: pip install sentence-transformers chromadb rank_bm25
"""
from __future__ import annotations
from collections import defaultdict
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi
import chromadb
 
 
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
RERANK_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"  # tiny + strong
COLLECTION = "docs_v1"
K_RETRIEVE, K_RERANK, K_FINAL = 50, 20, 5
 
 
class HybridRetriever:
    def __init__(self):
        self.client = chromadb.PersistentClient(path="./chroma_store")
        self.coll = self.client.get_collection(COLLECTION)
        self.embed = SentenceTransformer(EMBED_MODEL)
        self.rerank = CrossEncoder(RERANK_MODEL)
 
        # Load all docs once for BM25 (fine up to a few hundred K chunks).
        all_docs = self.coll.get(include=["documents", "metadatas"])
        self.docs = all_docs["documents"]
        self.ids = all_docs["ids"]
        self.metas = all_docs["metadatas"]
        tokenised = [d.lower().split() for d in self.docs]
        self.bm25 = BM25Okapi(tokenised)
 
    # ---------- Dense retrieval ----------
    def dense(self, q: str, k: int = K_RETRIEVE) -> list[tuple[str, str, dict]]:
        qv = self.embed.encode([q])[0].tolist()
        hits = self.coll.query(query_embeddings=[qv], n_results=k,
                               include=["documents", "metadatas"])
        return list(zip(hits["ids"][0], hits["documents"][0], hits["metadatas"][0]))
 
    # ---------- Sparse retrieval (BM25 in-memory) ----------
    def sparse(self, q: str, k: int = K_RETRIEVE) -> list[tuple[str, str, dict]]:
        scores = self.bm25.get_scores(q.lower().split())
        top = sorted(range(len(scores)), key=lambda i: -scores[i])[:k]
        return [(self.ids[i], self.docs[i], self.metas[i]) for i in top]
 
    # ---------- Reciprocal Rank Fusion ----------
    def fuse(self, *result_lists, k_const: int = 60) -> list[tuple[str, str, dict]]:
        score_by_id: dict[str, float] = defaultdict(float)
        doc_by_id: dict[str, tuple[str, dict]] = {}
        for results in result_lists:
            for rank, (id_, doc, meta) in enumerate(results, start=1):
                score_by_id[id_] += 1.0 / (k_const + rank)
                doc_by_id[id_] = (doc, meta)
        ranked_ids = sorted(score_by_id, key=lambda i: -score_by_id[i])
        return [(id_, *doc_by_id[id_]) for id_ in ranked_ids]
 
    # ---------- Cross-encoder rerank ----------
    def rerank_top(self, q: str, candidates, k: int = K_FINAL):
        pairs = [(q, doc) for _, doc, _ in candidates[:K_RERANK]]
        scores = self.rerank.predict(pairs)
        ranked = sorted(zip(scores, candidates[:K_RERANK]),
                        key=lambda x: -x[0])
        return [(s, id_, doc, meta) for s, (id_, doc, meta) in ranked[:k]]
 
    # ---------- One-shot: dense + sparse + fuse + rerank ----------
    def search(self, q: str) -> list[tuple[float, str, str, dict]]:
        dense_hits  = self.dense(q)
        sparse_hits = self.sparse(q)
        fused = self.fuse(dense_hits, sparse_hits)
        return self.rerank_top(q, fused, k=K_FINAL)
 
 
if __name__ == "__main__":
    r = HybridRetriever()
    for q in ("How do I refresh an OAuth token?",
              "error code E-4821 damaged shipment",
              "What is chunking in RAG?"):
        print(f"\n=== {q!r} ===")
        for score, id_, doc, meta in r.search(q):
            print(f"  [{score:+.3f}] {meta.get('breadcrumb','(no bcrumb)')}")
            print(f"          {doc[:120].replace(chr(10),' ')}...")