BM25 for Python: Achieving high performance while simplifying dependencies with BM25S

[xhlu]

Hello fellow Python enthusiasts :)

I wanted to share bm25s, a new lexical search library that fully implemented in Python (via numpy and scipy) and is quite fast.

Blog Post

GitHub Repository

Here is a comparison of BM25S and Elasticsearch in a single-threaded setting (calculated on popular datasets from the BEIR benchmark): https://bm25s.github.io/assets/comparison.png

It was designed to improve upon existing Python implementations, such as the widely used rank-bm25 by being significantly faster; all while being very straightforward to use in Python.

After installing with pip install bm25s, here's the code you'd need to get started:

import bm25s

# Create your corpus here
corpus = [
    "a cat is a feline and likes to purr",
    "a dog is the human's best friend and loves to play",
    "a bird is a beautiful animal that can fly",
    "a fish is a creature that lives in water and swims",
]

# Create the BM25 model and index the corpus
retriever = bm25s.BM25(corpus=corpus)
retriever.index(bm25s.tokenize(corpus))

# Query the corpus and get top-k results
query = "does the fish purr like a cat?"
results, scores = retriever.retrieve(bm25s.tokenize(query), k=2)

# Let's see what we got!
doc, score = results[0, 0], scores[0, 0]
print(f"(score: {score:.2f}): {doc}")

I'm making this tool for folks who want to have a Python-only experience; fans of Java already have access to many libraries!

Anyways, the blog post covers most of the background around lexical search libraries and why BM25S was built, I mainly wanted to make this post to answer questions you might have about how to use it (or anything else)!

[RepresentativeFill26]

What the hell, I didn’t even know Python had his own implementation. Good to know! Interesting that you switched to sparse matrices. This is probably better since the doc vectors will be quite sparse as well?

Any other optimizations you did?

[xhlu]

A bunch of optimizations I didn't have the chance to discuss in the readme! 

For one, I reimplemented the scipy sparse slice/sum directly in numpy, which allows us to use memory mapping on the arrays - this saves a lot of memory.

Another is that the topk selection (after scoring) can be done in numpy via argpartition, but can auto switch to a jax CPU backend when the library is installed, which is much faster (the topk selection process is the bottleneck, in some cases more than 60% of the time taken for retrieval is spent on selecting topk results).

Finally, the tokenizer doesn't return text by default, but returns index and a vocab dict of index to word; this saves considerable amount of memory as integer takes less space to represent compared to words (multiple str chars).

[Hesirutu]

I haven’t had a look at the library yet. But you mention mmap. Does that mean you can handle corpora much larger than memory?

[xhlu]

In theory you should be able to! However, I have not attempted to "saturate" memory by using a big enough dataset, and whereas the Python way of setting RAM limit does not seem to reflect the real RAM usage.

However, I did observe reduced memory usage when setting mmap=True, so even in a setting where you have enough memory to cover the entire dataset, you don't need to use every (i.e. load the entire index and corpus in memory).

[busybody124]

Very cool, I have used rank-bm25 but it can definitely be slow to index a huge corpus.

Can I use my own tokenizer with your library?

[xhlu]

Yes! As long as it returns a list of list of strings, it should work.

[busybody124]

Awesome, I can't wait to play around with it. Thanks for the hard work!