RAG Retrieval: BM25 + Embedding Rerank - Representation Learning & Retrieval

Hybrid Retrieval for RAG: BM25 + Embedding Rerank

What you'll build

A two-stage retriever: score every document with BM25, take the top-n as candidates, then rerank them with a weighted fusion of normalized BM25 and embedding cosine similarity.

Background

A retrieval-augmented generation (RAG) system answers a question by pasting a few retrieved passages into the LLM's prompt, so the answer can only be as good as the retrieval. Lexical search (BM25) finds rare exact terms that embeddings blur (names, error codes, IDs); dense search finds paraphrases that share no words. Production systems combine them: a cheap first stage over the whole corpus, then a more expensive or more semantic second stage over a few candidates (here a score fusion; often a cross-encoder).

Stage 1: BM25 (Robertson & Zaragoza)

For a query q and a document d (both token lists):

score(q, d) = sum over unique terms t of q:
              IDF(t) * f(t, d) * (k1 + 1) / ( f(t, d) + k1 * (1 - b + b * |d| / avgdl) )

IDF(t) = ln( (N - n_t + 0.5) / (n_t + 0.5) + 1 )

• f(t, d): count of t in d; |d|: number of tokens in d; avgdl: mean document length
• N: number of documents; n_t: number of documents containing t
• k1 (default 1.5) saturates repeated terms; b (default 0.75) normalizes for length
• The + 1 inside the log (Lucene's variant) keeps IDF positive even for very common terms

Stage 2: candidates

The top-n documents by BM25, descending; ties go to the lower document index.

Stage 3: rerank by weighted score fusion

For each candidate c, with cos_c the cosine similarity between query_emb and doc_embs[c] (denominator ||q|| · ||d|| + 1e-12):

bm25_norm_c = bm25_c / max over candidates of bm25     (all 0 if that max is 0)
fused_c     = alpha * bm25_norm_c + (1 - alpha) * cos_c

Sort the candidates by fused, descending; ties keep the BM25 candidate order. Dividing by the max puts BM25 on a [0, 1] scale comparable with cosine.

Inputs

• query_tokens: list of str; docs_tokens: list of N lists of str (no stemming: 'cats' != 'cat')
• query_emb: (D,); doc_embs: (N, D)
• n: number of candidates (default 3); alpha in [0, 1] (default 0.5); k1, b as above

Expected Return Format

• bm25_scores(query_tokens, docs_tokens, k1=1.5, b=0.75): (N,) float array
• retrieve_and_rerank(...): a dictionary with
  • 'bm25': (N,) BM25 scores of every document
  • 'candidates': (n,) integer array, top-n document indices by BM25
  • 'reranked': (n,) the same documents in fused order
  • 'fused': (n,) fused scores, aligned with 'reranked'

Edge cases and conventions

• Each distinct query term counts once (repeating a word in the query doesn't double its weight).
• alpha = 1 keeps the BM25 order; alpha = 0 sorts the candidates by cosine alone.

Constraints

NumPy and plain Python; no retrieval libraries.

Follow-up

• Reciprocal rank fusion scores sum over rankers of 1 / (60 + rank) and ignores raw scores. When is it more robust than weighted fusion?
• How would you measure the retriever on its own (recall@k, MRR) before blaming the LLM for a wrong answer?

Input: query = ['cat', 'pet'], docs = ['the cat sat on the mat', 'dogs and cats are pets', 'the cat chased the mouse', 'stock markets fell sharply today', 'a cat is a small pet'] (split on spaces)
Output: bm25 = [0.5133, 0.0, 0.5576, 0.0, 1.8336], candidates = [4, 2, 0]

Input: same query and docs, query_emb = [1.0, 0.5, 0.0], doc_embs rows 0, 2, 4 = [1, 0.5, 0], [0.6, 0.6, 0.4], [0, 0.3, 1]; n=3, alpha=0.5
Output: reranked = [0, 2, 4], fused = [0.64, 0.5811, 0.5643]

import numpy as np

def bm25_scores(query_tokens, docs_tokens, k1=1.5, b=0.75):
    """
    BM25 score of every document for one query.

    Args:
        query_tokens: list of query tokens (each distinct term counts once)
        docs_tokens: list of N token lists
        k1, b: BM25 parameters

    Returns:
        (N,) array of scores
    """
    # Your code here
    pass


def retrieve_and_rerank(query_tokens, docs_tokens, query_emb, doc_embs,
                        n=3, alpha=0.5, k1=1.5, b=0.75):
    """
    BM25 top-n, then rerank by alpha * (bm25 / max candidate bm25)
    + (1 - alpha) * cosine(query_emb, doc_emb).

    Args:
        query_emb: (D,) query embedding
        doc_embs: (N, D) document embeddings

    Returns:
        dict with 'bm25' (N,), 'candidates' (n,), 'reranked' (n,), 'fused' (n,)
    """
    # Your code here
    pass