BM25 › Implementing BM25 from Scratch

# Sample corpus about programming languages
corpus = [
    "Python is a high-level programming language known for its simple syntax and readability",
    "Java is a popular object-oriented programming language used for enterprise applications",
    "JavaScript is the programming language of the web enabling interactive web pages",
    "Python and JavaScript are both interpreted languages with dynamic typing",
    "Java requires compilation to bytecode which runs on the Java Virtual Machine",
    "The Python programming language supports multiple programming paradigms",
    "Web developers often use JavaScript frameworks like React and Vue",
    "Python is widely used for machine learning and data science applications",
    "Java enterprise applications often use Spring framework for dependency injection",
    "JavaScript can run on servers using Node.js runtime environment",
]

import re
from collections import defaultdict


class BM25Index:
    def __init__(self, k1=1.2, b=0.75):
        self.k1 = k1
        self.b = b
        self.documents = []
        self.doc_lengths = []
        self.avgdl = 0
        self.doc_freq = defaultdict(int)
        self.inverted_index = defaultdict(list)  # term -> [(doc_id, tf), ...]

    def tokenize(self, text):
        """Simple tokenization: lowercase and split on non-alphanumeric."""
        return re.findall(r"\b[a-z]+\b", text.lower())

    def index(self, documents):
        """Build the BM25 index from a list of documents."""
        self.documents = documents

        for doc_id, doc in enumerate(documents):
            tokens = self.tokenize(doc)
            self.doc_lengths.append(len(tokens))

            # Count term frequencies in this document
            term_freq = Counter(tokens)

            # Update document frequencies and inverted index
            for term, tf in term_freq.items():
                if tf > 0:
                    self.doc_freq[term] += 1
                    self.inverted_index[term].append((doc_id, tf))

        self.avgdl = sum(self.doc_lengths) / len(self.doc_lengths)
        self.N = len(documents)

def compute_idf(self, term):
    """Compute IDF for a term."""
    n = self.doc_freq.get(term, 0)
    return math.log((self.N - n + 0.5) / (n + 0.5))


def score_document(self, doc_id, query_terms):
    """Compute BM25 score for a single document."""
    doc_len = self.doc_lengths[doc_id]
    tokens = self.tokenize(self.documents[doc_id])
    term_freq = Counter(tokens)

    score = 0
    for term in query_terms:
        if term not in term_freq:
            continue

        tf = term_freq[term]
        idf = self.compute_idf(term)

        numerator = tf * (self.k1 + 1)
        denominator = tf + self.k1 * (
            1 - self.b + self.b * (doc_len / self.avgdl)
        )
        score += idf * (numerator / denominator)

    return score


def search(self, query, top_k=5):
    """Search the index and return top-k results."""
    query_terms = self.tokenize(query)

    # Find candidate documents (those containing at least one query term)
    candidates = set()
    for term in query_terms:
        for doc_id, _ in self.inverted_index.get(term, []):
            candidates.add(doc_id)

    # Score candidates
    results = []
    for doc_id in candidates:
        score = self.score_document(doc_id, query_terms)
        if score > 0:
            results.append((doc_id, score))

    # Sort by score descending
    results.sort(key=lambda x: x[1], reverse=True)
    return results[:top_k]


# Add methods to the class
BM25Index.compute_idf = compute_idf
BM25Index.score_document = score_document
BM25Index.search = search

# Build the index
index = BM25Index(k1=1.2, b=0.75)
index.index(corpus)

# Search for different queries
queries = [
    "python programming",
    "web javascript",
    "machine learning",
    "java enterprise",
]

search_results = {}
for query in queries:
    results = index.search(query, top_k=3)
    search_results[query] = results