Implementation of Okapi BM25 in python

Question

I am trying to implement Okapi BM25 in python. While I have seen some tutorials how to do it, it seems I am stuck in the process.

So I have collection of documents (and has as columns 'id' and 'text') and queries (and has as columns 'id' and 'text'). I have done the pre-processing steps and I have my documents and queries as a list:

documents = list(train_docs['text'])        #put the documents text to list
queries = list(train_queries_all['text'])   #put the queries text to list

Then for BM25 I do this:

pip install rank_bm25

#calculate BM25

from rank_bm25 import BM25Okapi

bm25 = BM25Okapi(documents)

#compute the score

bm_score = BM25Okapi.get_scores(documents, query=queries)

But it wouldn't work.

---

Then I tried to do this:

import math
import numpy as np
from multiprocessing import Pool, cpu_count

nd = len(documents) # corpus_size = 3612 (I am not sure if this is necessary)

class BM25:
    def __init__(self, documents, tokenizer=None):
        self.corpus_size = len(documents)
        self.avgdl = 0
        self.doc_freqs = []
        self.idf = {}
        self.doc_len = []
        self.tokenizer = tokenizer

        if tokenizer:
            documents = self._tokenize_corpus(documents)

        nd = self._initialize(documents)
        self._calc_idf(nd)

    def _initialize(self, documents):
        nd = {}  # word -> number of documents with word
        num_doc = 0
        for document in documents:
            self.doc_len.append(len(document))
            num_doc += len(document)

            frequencies = {}
            for word in document:
                if word not in frequencies:
                    frequencies[word] = 0
                frequencies[word] += 1
            self.doc_freqs.append(frequencies)

            for word, freq in frequencies.items():
                if word not in nd:
                    nd[word] = 0
                nd[word] += 1

        self.avgdl = num_doc / self.corpus_size
        return nd

    def _tokenize_corpus(self, documents):
        pool = Pool(cpu_count())
        tokenized_corpus = pool.map(self.tokenizer, documents)
        return tokenized_corpus

    def _calc_idf(self, nd):
        raise NotImplementedError()

    def get_scores(self, queries):
        raise NotImplementedError()

    def get_top_n(self, queries, documents, n=5):

        assert self.corpus_size == len(documents), "The documents given don't match the index corpus!"

        scores = self.get_scores(queries)
        top_n = np.argsort(scores)[::-1][:n]
        return [documents[i] for i in top_n]

class BM25T(BM25):
    def __init__(self, documents, k1=1.5, b=0.75, delta=1):
        # Algorithm specific parameters
        self.k1 = k1
        self.b = b
        self.delta = delta
        super().__init__(documents)

    def _calc_idf(self, nd):
        for word, freq in nd.items():
            idf = math.log((self.corpus_size + 1) / freq)
            self.idf[word] = idf

    def get_scores(self, queries):
        score = np.zeros(self.corpus_size)
        doc_len = np.array(self.doc_len)
        for q in queries:
            q_freq = np.array([(doc.get(q) or 0) for doc in self.doc_freqs])
            score += (self.idf.get(q) or 0) * (self.delta + (q_freq * (self.k1 + 1)) /
                                               (self.k1 * (1 - self.b + self.b * doc_len / self.avgdl) + q_freq))
        return score

and then I try to get the scores:

score = BM25.get_scores(self=documents, queries)

But I get as a meesage: score = BM25.get_scores(self=documents, queries)

SyntaxError: positional argument follows keyword argument

---

Does anyone has an idea why there is this error? Thank you in advance.

Answer 1

1 ) tokenize corpus or send tokinizing function to class

2 ) send only queries to "get_scores" function

read official example

from rank_bm25 import BM25Okapi

corpus = [
    "Hello there good man!",
    "It is quite windy in London",
    "How is the weather today?"
]

tokenized_corpus = [doc.split(" ") for doc in corpus]

bm25 = BM25Okapi(tokenized_corpus)

query = "windy London"
tokenized_query = query.split(" ")

doc_scores = bm25.get_scores(tokenized_query)

Answer 2

TLDR: Libraries suggested in other answers do not use the original Bm25 algorithm, nor does Wikipedia, go to the bottom if you want to run the original BM25 implementation.

I know there's already Birol's answer that concise show how to properly use Rank-BM25 as a library.

However, to answer your original question of "How to use Okapi BM25 in Python?" (paraphrasing from "implement" as it seems your goal is not to reimplement it), I believe it is useful to decompose the question to:

1. What is Okapi in Okapi BM25?
2. What is the original BM25 and how do I implement it?
3. Which other variants can I use?

For (1), the answer is that it is a platform/system developed by the City University of London's department of Information Science, built from 12 years of research done at the university (read more on their website). So it is NOT an algorithm, but rather a testbed for experiments!

(2) Note that Dr. Susan Jones, the author of the webpage, is also one of the author of the report, "OKAPI at TREC", one of the early published work about the Best Match family of algorithms (including BM25). Although the term frequency component is clearly given, it is unclear by directly reading the paper what algorithm they used. However, it is clear that when people refer to Okapi BM25, they mean this one!

(3) However, a later survey by Kamphuis et al. (2020) clearly explains the formula used in the original BM25 (i.e. the one by Robertson et al. from City University), as well as 7 other popular variants:

Note however that Lucene is likely what a lot of people are using when they are using BM25 through a library like Pyserini or Elastic. However, that is different from rank_bm25, which uses the inverse document frequency (IDF) from the original Robertson formulation, and the term frequency component (TFC) from the ATIRE paper.

To make things even more confusing, Wikipedia shows an implementation of BM25 that is also different from rank_bm25 and Lucene, as it uses the IDF from Lucene and the TFC from ATIRE. If you want to confirm that, you can take a look at the Okapi at TREC-3 paper where eq. 1 reduces to the IDF shown in the picture if you set R=0 and r=0, and eq. 2 matches exactly.

So does that mean there are NO true BM25 implementation?

Well, actually, Chris Kamphuis released SQL code for reproducing the survey linked above, but it is not in Python; and other libraries like rank_bm25 provides their own custom version of BM25.

Fortunately, the new bm25s library allows you to also run 5 of the popular variants of BM25, including the original one method="robertson", the popular Lucene implementation method="lucene" and 3 others (disclaimer: I am the original author).

• GitHub: github.com/xhluca/bm25s
• Blog Post: huggingface.co/blog/xhluca/bm25s
• Website: bm25s.github.io

To use it, simply install with pip install bm25s and run:

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",
]

# Tokenize the corpus and index it
corpus_tokens = bm25s.tokenize(corpus)
# Change method to use your preferred version
retriever = bm25s.BM25(corpus=corpus, method="robertson")
retriever.index(corpus_tokens)

# You can now search the corpus with a query
query = "does the fish purr like a cat?"
query_tokens = bm25s.tokenize(query)
docs, scores = retriever.retrieve(query_tokens, k=2)
print(f"Best result (score: {scores[0, 0]:.2f}): {docs[0, 0]}")

# Happy with your index? Save it for later...
retriever.save("bm25s_index_animals")

# ...and load it when needed
ret_loaded = bm25s.BM25.load("bm25s_index_animals", load_corpus=True)

Notes:

• The tokenizer is basic and supports English out of the box (but can be configured to use with other languages)
• The variants are: Robertson (original), Lucene (default), ATIRE, BM25L, BM25+
• Beyond supporting the five variants, bm25s is designed in a way where the relevance scores are calculated during indexing rather than during retrieval, which can lead to up to 100-500x faster retrieval on popular datasets from BEIR (more details on the website, bm25s.github.io)