BM25 scores identical to rank_bm25, up to 950× faster

Change one import

BM25Okapi, BM25L and BM25Plus take the same arguments as in rank_bm25 0.2.2, including tokenizer, and have the same methods and attributes:

from rank_bm25_rs import BM25Okapi  # was: from rank_bm25 import BM25Okapi

corpus = ["Hello there good man!", "It is quite windy in London", "How is the weather today?"]
bm25 = BM25Okapi([doc.split(" ") for doc in corpus])

bm25.get_scores(["windy", "London"])
# array([0.        , 0.93729472, 0.        ])
bm25.get_top_n(["windy", "London"], corpus, n=1)
# ['It is quite windy in London']

get_batch_scores, idf, doc_freqs, doc_len, avgdl and average_idf work as before, and indexes can be pickled. get_scores_many scores a list of queries on all CPU cores without the GIL and returns one row per query; RAYON_NUM_THREADS limits the threads.

Speed up libraries that import rank_bm25

LangChain's BM25Retriever and crawl4ai import rank_bm25 themselves. Call install() before importing them:

import rank_bm25_rs

rank_bm25_rs.install()  # `import rank_bm25` now returns rank_bm25_rs

from langchain_community.retrievers import BM25Retriever

Modules that imported rank_bm25 before the call keep the original.

How much faster

BM25Okapi on the documents of the CNN/DailyMail, XSum and PubMed test splits (lowercased, split on whitespace; 100,000 documents repeat the 29,482 available), queries of 6 tokens. Apple M1 Pro, CPython 3.13; every timed query was checked to give identical scores.

Method, raw results and the commands to reproduce them: BENCHMARKS.md.

Why it is so much faster, and still exact

For every query term, rank_bm25 builds a frequency array by looking the term up in a Python dictionary for every document of the corpus, then scores all documents with NumPy. The Python loop over the corpus is what costs 190 ms at 100,000 documents, even when the term occurs in a handful of them.

rank-bm25-rs keeps an inverted index: for each term, the documents that contain it and how often. For a document without the term, BM25Okapi and BM25L add exactly zero (or negative zero, which leaves the sum unchanged), so only the term's postings need visiting. BM25Plus adds the same nonzero amount to every document, which is one pass over a float array in Rust. Every operation that remains follows NumPy's order in rank_bm25's expressions, for example (1 - b) + (b * doc_len) / avgdl before multiplying by k1, so the floating-point results are the same bit for bit.

Where that shortcut would not hold, the index switches to a dense loop: when k1 = 0 or empty documents make the denominators zero, rank_bm25 returns NaN for every document, and so does rank-bm25-rs.

How "identical" is checked

The tests compare every score with rank_bm25 0.2.2 on CPython 3.10 to 3.15: corpora from one to thousands of documents, empty documents, unknown and repeated query terms, property-based tests and parameters that give NaN or negative scores. A separate workflow runs the full CNN/DailyMail, XSum and PubMed test splits with all three variants on Linux, macOS and Windows.

For anything the Rust index cannot reproduce exactly, rank_bm25's own code, bundled in the package, runs instead, so results, NumPy warnings and errors do not change:

• a corpus or query that is not a list or tuple of str, such as a generator or integer tokens;
• parameters that are NumPy scalars, NaN or infinite, or integers large enough to overflow NumPy's integer arithmetic;
• k1, b, delta or avgdl changed after construction, subclasses that override how the index is built, and non-default np.errstate;
• queries whose scores are NaN or infinite.

Two known differences remain: changes made to the idf, doc_freqs or doc_len attributes after construction are not seen by the index, and Python 3.10 or newer is required.

Rust

use rank_bm25::{Builder, Params};

let mut builder = Builder::new(Params::okapi());
for doc in ["Hello there good man!", "It is quite windy in London", "How is the weather today?"] {
    builder.add_document(doc.split(' '))?;
}
let index = builder.build()?;
let scores = index.get_scores(&["windy", "London"]);

Params::l() and Params::plus() select the other variants. Index is Send + Sync and can be serialised with to_bytes.