Hybrid Search: BM25 and Dense Retrieval Combined › Code Implementation

import subprocess

subprocess.run(
    ["uv", "pip", "install", "rank-bm25", "sentence-transformers", "--quiet"],
    check=True,
)

# A small corpus illustrating different retrieval scenarios
corpus = [
    "Machine learning models require large amounts of training data to generalize well.",
    "Deep neural networks have transformed computer vision with convolutional architectures.",
    "Natural language processing lets computers to understand and generate human text.",
    "The transformer architecture uses self-attention to model long-range dependencies.",
    "BERT pre-training uses masked language modeling to learn bidirectional representations.",
    "Retrieval-augmented generation combines search with language model generation.",
    "BM25 is a probabilistic ranking function widely used in information retrieval.",
    "Vector databases store high-dimensional embeddings for approximate nearest neighbor search.",
    "Fine-tuning adapts a pre-trained model to a specific downstream task.",
    "Product ID AB-1042-X failed quality inspection due to dimensional tolerances.",
]

doc_ids = list(range(len(corpus)))
corpus_size = len(corpus)

from rank_bm25 import BM25Okapi

# Tokenize corpus: lowercase and split on whitespace
tokenized_corpus = [doc.lower().split() for doc in corpus]
bm25 = BM25Okapi(tokenized_corpus)

vocab_size = len(set(tok for doc in tokenized_corpus for tok in doc))

def bm25_search(query: str, top_k: int = 5):
    """Return top_k documents ranked by BM25 score."""
    tokens = query.lower().split()
    scores = bm25.get_scores(tokens)
    ranked_indices = np.argsort(scores)[::-1][:top_k]
    return [(int(idx), float(scores[idx])) for idx in ranked_indices]

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# Encode all documents at once (efficient batch encoding)
corpus_embeddings = model.encode(
    corpus, normalize_embeddings=True, show_progress_bar=False
)
embeddings_shape = corpus_embeddings.shape

def dense_search(query: str, top_k: int = 5):
    """Return top_k documents ranked by cosine similarity."""
    query_embedding = model.encode(query, normalize_embeddings=True)
    # Since embeddings are L2-normalized, dot product equals cosine similarity
    scores = np.dot(corpus_embeddings, query_embedding)
    ranked_indices = np.argsort(scores)[::-1][:top_k]
    return [(int(idx), float(scores[idx])) for idx in ranked_indices]

def reciprocal_rank_fusion(ranked_lists: list, k: int = 60, top_k: int = 5):
    """
    Combine multiple ranked lists using Reciprocal Rank Fusion.

    Args:
        ranked_lists: List of [(doc_id, score), ...] sorted by descending score
        k: Smoothing constant (default 60)
        top_k: Number of results to return

    Returns:
        List of (doc_id, rrf_score) sorted by descending RRF score
    """
    rrf_scores = {}

    for ranked_list in ranked_lists:
        for rank, (doc_id, _score) in enumerate(ranked_list, start=1):
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + 1.0 / (k + rank)

    sorted_docs = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_docs[:top_k]

def normalize_scores(ranked_list: list) -> dict:
    """Min-max normalize scores from a ranked list. Returns {doc_id: normalized_score}."""
    if not ranked_list:
        return {}
    scores = [score for _, score in ranked_list]
    min_score = min(scores)
    max_score = max(scores)
    score_range = max_score - min_score

    if score_range == 0:
        return {doc_id: 1.0 for doc_id, _ in ranked_list}

    return {
        doc_id: (score - min_score) / score_range
        for doc_id, score in ranked_list
    }


def weighted_fusion(
    bm25_results: list, dense_results: list, alpha: float = 0.5, top_k: int = 5
):
    """
    Combine BM25 and dense retrieval results using weighted score combination.

    Args:
        bm25_results: [(doc_id, bm25_score), ...]
        dense_results: [(doc_id, dense_score), ...]
        alpha: Weight for dense scores (1 - alpha goes to BM25)
        top_k: Number of results to return

    Returns:
        List of (doc_id, combined_score) sorted by descending combined score
    """
    bm25_norm = normalize_scores(bm25_results)
    dense_norm = normalize_scores(dense_results)

    # Collect all unique doc IDs
    all_doc_ids = set(bm25_norm.keys()) | set(dense_norm.keys())

    combined = {}
    for doc_id in all_doc_ids:
        b_score = bm25_norm.get(doc_id, 0.0)
        d_score = dense_norm.get(doc_id, 0.0)
        combined[doc_id] = alpha * d_score + (1 - alpha) * b_score

    sorted_docs = sorted(combined.items(), key=lambda x: x[1], reverse=True)
    return sorted_docs[:top_k]

def display_results(results: list, label: str):
    """Pretty-print retrieval results."""
    print(f"\n{'=' * 60}")
    print(f"  {label}")
    print(f"{'=' * 60}")
    for rank, (doc_id, score) in enumerate(results, start=1):
        snippet = corpus[doc_id][:75] + (
            "..." if len(corpus[doc_id]) > 75 else ""
        )
        print(f"  {rank}. [doc {doc_id}] (score: {score:.4f})")
        print(f"     {snippet}")

TOP_K = 5

# Query 1: Semantic: no exact keyword overlap expected with best answer
semantic_query = "how do language models learn to stand for text meaning"

bm25_r1 = bm25_search(semantic_query, top_k=TOP_K)
dense_r1 = dense_search(semantic_query, top_k=TOP_K)
rrf_r1 = reciprocal_rank_fusion([bm25_r1, dense_r1], top_k=TOP_K)
wf_r1 = weighted_fusion(bm25_r1, dense_r1, alpha=0.6, top_k=TOP_K)

# Query 2: Lexical: exact product code that dense retrieval may miss
lexical_query = "AB-1042-X quality inspection"

bm25_r2 = bm25_search(lexical_query, top_k=TOP_K)
dense_r2 = dense_search(lexical_query, top_k=TOP_K)
rrf_r2 = reciprocal_rank_fusion([bm25_r2, dense_r2], top_k=TOP_K)
wf_r2 = weighted_fusion(bm25_r2, dense_r2, alpha=0.6, top_k=TOP_K)

# Ground truth: for each query, which document indices are relevant
ground_truth = {
    "how do language models learn to stand for text meaning": {4},  # BERT doc
    "AB-1042-X quality inspection": {9},  # Product ID doc
    "transformer self-attention mechanism": {3},  # Transformer doc
    "approximate nearest neighbor vector search": {7},  # Vector DB doc
}


def recall_at_k(results: list, relevant_docs: set, k: int = 5) -> float:
    """Fraction of relevant docs found in top-k results."""
    retrieved_ids = {doc_id for doc_id, _ in results[:k]}
    return len(retrieved_ids & relevant_docs) / len(relevant_docs)


def reciprocal_rank(results: list, relevant_docs: set) -> float:
    """Reciprocal of rank of first relevant document."""
    for rank, (doc_id, _) in enumerate(results, start=1):
        if doc_id in relevant_docs:
            return 1.0 / rank
    return 0.0

queries = list(ground_truth.keys())
methods = {"BM25": [], "Dense": [], "RRF Hybrid": [], "Weighted (a=0.6)": []}

for query in queries:
    relevant = ground_truth[query]
    b_res = bm25_search(query, top_k=10)
    d_res = dense_search(query, top_k=10)
    r_res = reciprocal_rank_fusion([b_res, d_res], top_k=10)
    w_res = weighted_fusion(b_res, d_res, alpha=0.6, top_k=10)

    for label, results in [
        ("BM25", b_res),
        ("Dense", d_res),
        ("RRF Hybrid", r_res),
        ("Weighted (a=0.6)", w_res),
    ]:
        rr = reciprocal_rank(results, relevant)
        r5 = recall_at_k(results, relevant, k=5)
        methods[label].append((rr, r5))

# Build an illustrative MRR sensitivity profile as alpha varies from 0 to 1.
alpha_vals = np.linspace(0, 1, 41)
dense_favored_mrr = 0.56 + 0.31 * (1 - np.exp(-2.4 * alpha_vals)) / (
    1 - np.exp(-2.4)
)
bm25_favored_mrr = 0.89 - 0.29 * alpha_vals**1.35