Problem 110 · Retrieval

BM25 Scoring from Scratch

MediumAI engineering

01Problem

Implement a `BM25` class over a list of documents with parameters `k1=1.5`, `b=0.75`. Tokenise by lowercasing and extracting `[a-z0-9]+`. Use the non-negative IDF `log(1 + (N - df + 0.5) / (df + 0.5))`. Provide `score(query, doc_index)` and `search(query, k)` returning `(doc_index, score)` pairs, highest first, ties by lower index, omitting zero scores.

02Key insight

BM25 is TF-IDF with two corrections that matter: term frequency saturates (k1), so the tenth mention adds little, and scores are normalised by document length relative to the average (b), so long documents do not win just by containing more words. The `1 +` inside the log keeps IDF positive for terms in more than half the corpus; the original formula goes negative and can make a matching document rank below a non-matching one. Lexical search still beats embeddings on exact identifiers, error codes and rare names, which is why hybrid retrieval keeps it.

03Solution

solution.pypython38 lines

import math
import re
from collections import Counter

_TOKEN = re.compile(r"[a-z0-9]+")


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


class BM25:
    def __init__(self, docs: list[str], k1: float = 1.5, b: float = 0.75) -> None:
        self.k1, self.b = k1, b
        self.tfs = [Counter(tokenize(d)) for d in docs]
        self.lens = [sum(tf.values()) for tf in self.tfs]
        self.n = len(docs)
        self.avgdl = (sum(self.lens) / self.n) if self.n else 0.0
        df: Counter[str] = Counter()
        for tf in self.tfs:
            df.update(tf.keys())
        self.idf = {t: math.log(1 + (self.n - c + 0.5) / (c + 0.5)) for t, c in df.items()}

    def score(self, query: str, i: int) -> float:
        tf, dl = self.tfs[i], self.lens[i]
        norm = self.k1 * (1 - self.b + self.b * dl / self.avgdl) if self.avgdl else self.k1
        total = 0.0
        for term in set(tokenize(query)):
            f = tf.get(term, 0)
            if f:
                total += self.idf[term] * f * (self.k1 + 1) / (f + norm)
        return total

    def search(self, query: str, k: int) -> list[tuple[int, float]]:
        scored = [(i, self.score(query, i)) for i in range(self.n)]
        scored = [s for s in scored if s[1] > 0]
        scored.sort(key=lambda s: (-s[1], s[0]))
        return scored[:k]

Try it first.

TimeIndex O(total tokens); query O(N * |q|) here, O(sum of posting lengths) with an inverted index

SpaceO(total tokens)

tests.py · 8 passing

tests.pypython20 lines

docs = [
    "the cat sat on the mat",
    "the dog sat on the log",
    "cat cat cat",
    "error code E1234 in payment service",
    "the cat sat on the mat while the dog slept on the long log by the door",
]
bm = BM25(docs)
assert bm.search("E1234", 3) == [(3, bm.score("E1234", 3))]
assert bm.idf["e1234"] > bm.idf["the"] > 0
assert bm.score("cat mat", 0) > bm.score("cat mat", 4)
assert bm.score("cat", 2) > bm.score("cat", 0)
assert bm.score("unicorn", 0) == 0 and bm.search("unicorn", 5) == []
flat = BM25(docs, b=0.0)
assert abs(flat.score("mat", 0) - flat.score("mat", 4)) < 1e-12
n, df, f, dl = 5, 2, 1, 6
avg = sum(bm.lens) / 5
expected = math.log(1 + (n - df + 0.5) / (df + 0.5)) * f * 2.5 / (f + 1.5 * (0.25 + 0.75 * dl / avg))
assert abs(bm.score("dog", 1) - expected) < 1e-12
assert [i for i, _ in bm.search("cat", 3)] == [2, 0, 4]

04Staff-level follow-up

Replace the per-document scan with an inverted index and WAND or block-max WAND so you skip documents that cannot enter the top-k. Then discuss tokenisation for code and non-English text, field boosts (title vs body), and why BM25 parameters tuned on one corpus rarely transfer.