AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Implement BM25 keyword search from scratch

Implement BM25 Keyword Search from Scratch

Every production RAG system uses hybrid search — vector (semantic) retrieval combined with keyword (BM25) retrieval. Understanding BM25 from first principles lets you tune it for your corpus, debug unexpected ranking behavior, and explain to stakeholders why keyword search still matters in the era of embeddings.

BM25 improves over plain TF-IDF in two critical ways: term saturation (doubling occurrences should not double the score — diminishing returns) and document length normalization (5 occurrences in a 100-word document is more significant than 5 in a 10,000-word document).

The BM25 formula

BM25(d, Q) = sum_t [ IDF(t) * tf(t,d) * (k1+1) / (tf(t,d) + k1*(1 - b + b*|d|/avgdl)) ]
IDF(t) = log((N - df(t) + 0.5) / (df(t) + 0.5) + 1)

Standard defaults: k1=1.5 (term saturation), b=0.75 (length normalization).

What you are building

Implement a BM25Index class that:

Builds the index — tokenize documents (lowercase + whitespace split), compute IDF for each term.
Scores documents for a query using the full BM25 formula.
Returns top-K results sorted by score descending, excluding zero-score documents.

Constraints

Standard library only. k1 and b configurable at init time. Type-hint everything.

Retrieval / medium / Step 14 of 15

Practice stage

Retrieval quality and ranking

Hint

Do not trust a single score blindly. Combine ranking logic with metadata and think about why a candidate deserves to rise.

Next drill

Build a multi-stage RAG pipeline with query routing

Success criteria

- Ranking rule is explainable
- Supports future iteration without rewriting everything
- Feels tied to product trust rather than pure math trivia

Review checklist

- Did I account for both relevance and metadata?
- Would I be able to explain the ranking rule in an interview?
- Does the code make future tuning easy?

from __future__ import annotations

import math
from collections import Counter, defaultdict
from dataclasses import dataclass

@dataclass
class BM25Result:
    doc_id: str
    text: str
    score: float

def tokenize(text: str) -> list[str]:
    return text.lower().split()

class BM25Index:
    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self._docs: dict[str, str] = {}
        self._doc_tf: dict[str, Counter] = {}
        self._df: dict[str, int] = defaultdict(int)
        self._doc_lengths: dict[str, int] = {}
        self._avgdl: float = 0.0

def add_documents(self, docs: list[dict]) -> None:
        """Each doc: {\"id\": str, \"text\": str}. Tokenize, store TF, accumulate DF, compute avgdl."""
        # TODO: implement
        raise NotImplementedError

def _idf(self, term: str) -> float:
        """IDF = log((N - df + 0.5) / (df + 0.5) + 1)"""
        # TODO: implement
        raise NotImplementedError

def _bm25_score(self, doc_id: str, query_terms: list[str]) -> float:
        """Sum BM25 formula over all query terms for one document."""
        # TODO: implement
        raise NotImplementedError

def search(self, query: str, top_k: int = 10) -> list[BM25Result]:
        """Return top-k documents by BM25 score (exclude zero-score docs)."""
        # TODO: implement
        raise NotImplementedError

Practice

Generate a variation

Generate a new exercise variation to deepen understanding or practice a related concept.

Attempt history

Recent submissions

Before you submit, decide what a strong answer should make obvious to the reviewer.

No attempts yet.