BM25 from scratch in Python: 27 lines, score-identical to rank_bm25 and up to 1.7x faster

Can you replace rank_bm25 with a short BM25 you wrote yourself, and what do you give up? I wrote an inverted-index BM25 and compared it with rank_bm25 on the FastAPI docs.

Answer: with the same idf formula, the two agree to floating-point noise (max score difference 1.8e-13 over 500 queries × 1,504 chunks), with identical top-10 sets on all 500 queries. The inverted index was 1.7x faster on short queries and 1.2x faster on 8-token queries. The one thing that does change rankings is the idf formula, not the implementation.

If you want the fusion side of this, see hybrid search with RRF. For the basics, the free guide on retrieval covers why lexical search still matters.

Setup

• Corpus: FastAPI English docs split at Markdown headings (levels 1-4), keeping sections over 200 characters. That gives 1,504 chunks, average 137.9 tokens, vocabulary 9,838.
• Tokenizer: lowercase, \w+. Both implementations get identical token lists.
• Parameters: k1=1.5, b=0.75 for both.
• Queries: 300 section headings with at least 2 tokens, and 200 queries made of 8 random tokens from a random chunk.
• Checks: (1) my code with rank_bm25’s idf; (2) my code with the Lucene-style idf; (3) speed.

The code

import math, re
from collections import Counter, defaultdict

def tokenize(t): return re.findall(r"\w+", t.lower())

class BM25:
    def __init__(self, docs, k1=1.5, b=0.75):
        self.k1, self.b, self.N = k1, b, len(docs)
        self.dl = [len(d) for d in docs]
        self.avgdl = sum(self.dl) / self.N
        self.post = defaultdict(list)            # term -> [(doc_id, tf)]
        for i, d in enumerate(docs):
            for t, f in Counter(d).items():
                self.post[t].append((i, f))
        # Lucene-style idf: never negative
        self.idf = {t: math.log(1 + (self.N - len(p) + 0.5) / (len(p) + 0.5))
                    for t, p in self.post.items()}

    def scores(self, query):
        s = defaultdict(float)
        for t in query:
            idf = self.idf.get(t)
            if idf is None: continue
            for i, f in self.post[t]:
                norm = f + self.k1 * (1 - self.b + self.b * self.dl[i] / self.avgdl)
                s[i] += idf * f * (self.k1 + 1) / norm
        return s

    def top(self, query, k=10):
        s = self.scores(query)
        return sorted(s, key=lambda i: (-s[i], i))[:k]

Only documents that share a term with the query get a score, which is why it is faster than looping over every document per query term.

Results

Correctness, using rank_bm25’s idf (log((N-n+0.5)/(n+0.5)), with negative values replaced by 0.25 × the average idf, as in the library):

check	result
max absolute score difference	1.776e-13
top-10 set identical	500 / 500
top-1 identical	500 / 500

I compared only documents with a non-zero score. rank_bm25 returns every document, so zero-score filler would otherwise show up as a false mismatch (my first run did exactly that and reported 6 differences).

Lucene-style idf (log(1 + (N-n+0.5)/(n+0.5)), the smoothed form most search engines use) against rank_bm25:

check	result
mean top-10 overlap	0.895
top-1 identical	486 / 500

So 14 of 500 queries (2.8%) changed their best result. The likely reason: rank_bm25 handles terms that appear in more than half of the chunks by flooring a negative idf, while the Lucene form gives them a small positive weight. I didn’t isolate this by experiment, so treat it as the probable cause, not a measured one.

Speed (ms per query, best of 3 passes; rank_bm25 timing includes argsort for the top 10):

query set	rank_bm25	from scratch	speedup
heading (2+ tokens)	1.405	0.810	1.7x
8 random tokens	1.662	1.436	1.2x

Index build went the other way: 39 ms for rank_bm25, 58 ms for mine (1.5x slower).

What the numbers mean

The implementation is not the risk. BM25 is a short formula, and with matching idf the scores are the same to 13 decimal places. The decisions that move results are the ones around it: the idf variant, the tokenizer, k1 and b.

The speed gain comes from skipping documents that share no query term. Longer queries touch more posting lists, so the advantage shrinks. Neither library is slow at this scale; both finish in under 2 ms per query on 1,504 chunks.

Caveats

• One corpus of 1,504 chunks. At 100x the size, the inverted index should pull further ahead, but I didn’t measure it.
• I didn’t measure retrieval quality (MRR or recall) here, only agreement between implementations. Whether Lucene idf is better for your data needs your own golden set.
• Timings come from a shared 2-core CI runner, so differences under about 20% are within noise. The 1.2x result is the least reliable.
• Pure Python. For millions of documents use Lucene-based engines or a sparse-matrix library.
• Stemming, stop words and phrase queries are not covered.

Rule of thumb

Use rank_bm25 for prototypes, and write your own when you need a persistent index, a different idf, or per-query speed. Either way, check yours against the library with identical tokens and rank_bm25’s idf first. If scores differ by more than 1e-9, you have a bug, not a design choice.

Experiment files

bench_bm25.py and results.txt are in the Experiment files box below. The script needs the corpus from fetch_corpus.sh and runs in a few seconds.