Keyword Search with BM25

import re
from math import log
from collections import Counter

def tokenise(text: str) -> list[str]:
    return re.findall(r'[a-z0-9_]+', text.lower())

# IDF(t) = log((N - df + 0.5) / (df + 0.5) + 1)
# N  = total number of documents
# df = number of documents containing term t

class BM25Index:
    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.chunks: list[str] = []
        self.tf: list[Counter] = []
        self.idf: dict[str, float] = {}
        self.avg_dl: float = 0.0

    def add(self, chunks: list[str]) -> None:
        self.chunks = chunks
        tokenised = [tokenise(c) for c in chunks]
        self.tf = [Counter(t) for t in tokenised]
        N = len(chunks)
        self.avg_dl = sum(len(t) for t in tokenised) / max(N, 1)
        df: Counter = Counter()
        for tokens in tokenised:
            for term in set(tokens):
                df[term] += 1
        self.idf = {
            term: log((N - freq + 0.5) / (freq + 0.5) + 1)
            for term, freq in df.items()
        }

    def search(self, query: str, top_k: int = 3) -> list[tuple[float, str]]:
        q_tokens = tokenise(query)
        scores = []
        for i, chunk in enumerate(self.chunks):
            dl = sum(self.tf[i].values())
            score = 0.0
            for term in q_tokens:
                if term not in self.idf:
                    continue
                tf_val = self.tf[i].get(term, 0)
                num = tf_val * (self.k1 + 1)
                denom = tf_val + self.k1 * (1 - self.b + self.b * dl / self.avg_dl)
                score += self.idf[term] * num / denom
            scores.append((score, chunk))
        scores.sort(key=lambda x: x[0], reverse=True)
        return scores[:top_k]