BM25 for Developers: What Actually Matters in Production

from rank_bm25 import BM25Plusimport numpy as npdef evaluate_bm25(tokenized_corpus, queries, ground_truth, k1, b):    """    Returns Precision@5 and Recall@10 for this parameter combination    """    bm25 = BM25Plus(tokenized_corpus, k1=k1, b=b)        precisions = []    recalls = []        for query_id, query_tokens in enumerate(queries):        scores = bm25.get_scores(query_tokens)        top_5_indices = np.argsort(scores)[-5:][::-1]        top_10_indices = np.argsort(scores)[-10:][::-1]                relevant_docs = ground_truth[query_id]                # Precision@5: how many of top 5 are relevant?        p5 = len(set(top_5_indices) & relevant_docs) / 5        precisions.append(p5)                # Recall@10: what fraction of relevant docs are in top 10?        r10 = len(set(top_10_indices) & relevant_docs) / len(relevant_docs) if relevant_docs else 0        recalls.append(r10)        return np.mean(precisions), np.mean(recalls)# Grid searchk1_values = [0.8, 1.0, 1.2, 1.5, 1.8, 2.0]b_values = [0.3, 0.5, 0.75, 0.9]best_params = Nonebest_score = 0for k1 in k1_values:    for b in b_values:        p5, r10 = evaluate_bm25(tokenized_corpus, queries, ground_truth, k1, b)        # F1-like combination of precision and recall        score = 2 * p5 * r10 / (p5 + r10) if (p5 + r10) > 0 else 0                print(f"k1={k1}, b={b}: P@5={p5:.3f}, R@10={r10:.3f}, F1={score:.3f}")                if score > best_score:            best_score = score            best_params = (k1, b)print(f"\nBest: k1={best_params[0]}, b={best_params[1]}")

k1=1.2, b=0.75: P@5=0.68, R@10=0.72, F1=0.70k1=1.5, b=0.75: P@5=0.71, R@10=0.69, F1=0.70k1=1.8, b=0.75: P@5=0.69, R@10=0.71, F1=0.70k1=1.5, b=0.50: P@5=0.74, R@10=0.68, F1=0.71  ← Best

import nltkfrom nltk.corpus import stopwordsfrom nltk.stem import WordNetLemmatizerfrom nltk.tokenize import word_tokenizefrom rank_bm25 import BM25Plusimport numpy as npfrom typing import List, Tuple, Set# Download NLTK data oncetry:    nltk.data.find('tokenizers/punkt')except LookupError:    nltk.download('punkt', quiet=True)    nltk.download('stopwords', quiet=True)    nltk.download('wordnet', quiet=True)    nltk.download('averaged_perceptron_tagger', quiet=True)class BM25Search:    def __init__(self, k1=1.5, b=0.75, use_stopwords=True, use_lemmatization=True):        """        k1: term frequency saturation (1.2-2.0)        b: length normalization (0.0-1.0)        use_stopwords: whether to remove common words        use_lemmatization: whether to reduce words to root form        """        self.k1 = k1        self.b = b        self.use_stopwords = use_stopwords        self.use_lemmatization = use_lemmatization                self.stop_words = set(stopwords.words('english')) if use_stopwords else set()        self.lemmatizer = WordNetLemmatizer() if use_lemmatization else None                self.documents = []        self.bm25 = None            def preprocess(self, text: str) -> List[str]:        """        Tokenize, normalize, and clean text        """        # Lowercase and tokenize        tokens = word_tokenize(text.lower())                # Keep only alphanumeric        tokens = [t for t in tokens if t.isalnum()]                # Remove stopwords        if self.use_stopwords:            tokens = [t for t in tokens if t not in self.stop_words]                # Lemmatize        if self.use_lemmatization:            # Simple lemmatization without POS tagging for speed            tokens = [self.lemmatizer.lemmatize(t) for t in tokens]                return tokens        def index(self, documents: List[str]):        """        Index documents for BM25 search        """        self.documents = documents        tokenized = [self.preprocess(doc) for doc in documents]                # Filter out empty documents        valid_docs = [(i, tokens) for i, tokens in enumerate(tokenized) if tokens]                if not valid_docs:            raise ValueError("No valid documents after preprocessing")                self.doc_indices = [i for i, _ in valid_docs]        tokenized_valid = [tokens for _, tokens in valid_docs]                self.bm25 = BM25Plus(tokenized_valid, k1=self.k1, b=self.b)                # Compute stats        doc_lengths = [len(tokens) for tokens in tokenized_valid]        self.avg_doc_length = np.mean(doc_lengths)        self.min_doc_length = min(doc_lengths)        self.max_doc_length = max(doc_lengths)                print(f"Indexed {len(tokenized_valid)} documents")        print(f"Avg length: {self.avg_doc_length:.1f} tokens")        print(f"Length range: {self.min_doc_length}-{self.max_doc_length} tokens")        def search(self, query: str, top_k=10) -> List[Tuple[int, float, str]]:        """        Search for top-k documents matching query                Returns: List of (original_doc_index, score, document_text)        """        if self.bm25 is None:            raise ValueError("Must call index() before search()")                query_tokens = self.preprocess(query)                if not query_tokens:            # Query became empty after preprocessing            return []                scores = self.bm25.get_scores(query_tokens)                # Get top-k indices        top_indices = np.argsort(scores)[-top_k:][::-1]                results = []        for idx in top_indices:            original_idx = self.doc_indices[idx]            score = scores[idx]            doc = self.documents[original_idx]            results.append((original_idx, score, doc))                return results        def explain_score(self, query: str, doc_index: int) -> dict:        """        Explain why a document scored the way it did        Useful for debugging        """        query_tokens = self.preprocess(query)        doc_tokens = self.preprocess(self.documents[doc_index])                # Which query tokens appear in doc?        matching_tokens = set(query_tokens) & set(doc_tokens)                explanation = {            'query_tokens': query_tokens,            'doc_tokens': doc_tokens,            'matching_tokens': list(matching_tokens),            'doc_length': len(doc_tokens),            'avg_doc_length': self.avg_doc_length,            'length_ratio': len(doc_tokens) / self.avg_doc_length if self.avg_doc_length > 0 else 0        }                return explanation# Example usageif __name__ == "__main__":    corpus = [        "Python is a popular programming language for data science and AI.",        "Machine learning and deep learning are subsets of artificial intelligence.",        "The quick brown fox jumps over the lazy dog.",        "Developers use Python for natural language processing and search engines.",        "Dogs are loyal animals, often considered man's best friend."    ]        # Initialize with tuned parameters    search = BM25Search(k1=1.5, b=0.75)    search.index(corpus)        # Search    query = "python search ai"    results = search.search(query, top_k=3)        print(f"\nQuery: '{query}'\n")    for rank, (idx, score, doc) in enumerate(results, 1):        print(f"{rank}. [Score: {score:.2f}] {doc}")        # Explain top result    print("\n--- Explaining top result ---")    explanation = search.explain_score(query, results[0][0])    print(f"Query tokens: {explanation['query_tokens']}")    print(f"Matching tokens: {explanation['matching_tokens']}")    print(f"Doc length: {explanation['doc_length']} (avg: {explanation['avg_doc_length']:.1f})")

TF-IDF:     Precision@5 = 0.64, Recall@10 = 0.71BM25 (k1=1.5, b=0.75): Precision@5 = 0.71, Recall@10 = 0.74

from sentence_transformers import SentenceTransformerimport numpy as npclass HybridSearch:    def __init__(self, bm25_weight=0.5):        self.bm25 = BM25Search(k1=1.5, b=0.75)        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')        self.bm25_weight = bm25_weight        self.dense_weight = 1 - bm25_weight        self.doc_embeddings = None        def index(self, documents):        # Index for BM25        self.bm25.index(documents)                # Compute dense embeddings        print("Computing dense embeddings...")        self.doc_embeddings = self.encoder.encode(            documents,            show_progress_bar=True,            convert_to_numpy=True        )        def search(self, query, top_k=10):        # BM25 scores        bm25_results = self.bm25.search(query, top_k=len(self.bm25.documents))        bm25_scores = {idx: score for idx, score, _ in bm25_results}                # Dense scores (cosine similarity)        query_embedding = self.encoder.encode(query, convert_to_numpy=True)        dense_scores = np.dot(self.doc_embeddings, query_embedding)                # Normalize both to [0, 1]        def normalize(scores_dict):            scores = np.array(list(scores_dict.values()))            if scores.max() == scores.min():                return {k: 0.5 for k in scores_dict}            return {                k: (v - scores.min()) / (scores.max() - scores.min())                for k, v in scores_dict.items()            }                norm_bm25 = normalize(bm25_scores)                dense_dict = {i: score for i, score in enumerate(dense_scores)}        norm_dense = normalize(dense_dict)                # Combine        final_scores = {}        for idx in range(len(self.bm25.documents)):            bm25_s = norm_bm25.get(idx, 0)            dense_s = norm_dense.get(idx, 0)            final_scores[idx] = self.bm25_weight * bm25_s + self.dense_weight * dense_s                # Sort and return top-k        ranked = sorted(final_scores.items(), key=lambda x: x[1], reverse=True)[:top_k]                return [            (idx, score, self.bm25.documents[idx])            for idx, score in ranked        ]