BM25 › BM25 vs TF-IDF: An Empirical Comparison

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Larger test corpus
test_corpus = [
    "Machine learning algorithms learn patterns from data",
    "Deep learning uses neural networks with many layers",
    "Natural language processing enables computers to understand text",
    "Computer vision allows machines to interpret images",
    "Reinforcement learning trains agents through rewards and penalties",
    "Supervised learning requires labeled training examples",
    "Unsupervised learning finds patterns without labels",
    "Transfer learning reuses knowledge from one task to another",
    "Neural networks are inspired by biological neurons",
    "Gradient descent optimizes model parameters iteratively",
    "Backpropagation computes gradients for neural network training",
    "Convolutional networks excel at image recognition tasks",
    "Recurrent networks process sequential data like text",
    "Attention mechanisms help models focus on relevant inputs",
    "Transformers revolutionized natural language processing",
    "BERT provides contextual word embeddings for NLP tasks",
    "GPT generates coherent text using autoregressive modeling",
    "Word embeddings represent words as dense vectors",
    "Sentiment analysis determines the emotional tone of text",
    "Named entity recognition identifies people places and organizations",
]

# Build BM25 index
bm25_index = BM25Index(k1=1.2, b=0.75)
bm25_index.index(test_corpus)

# Build TF-IDF index
tfidf_vectorizer = TfidfVectorizer()
tfidf_matrix = tfidf_vectorizer.fit_transform(test_corpus)


def tfidf_search(query, top_k=5):
    """Search using TF-IDF cosine similarity."""
    query_vec = tfidf_vectorizer.transform([query])
    similarities = cosine_similarity(query_vec, tfidf_matrix)[0]
    top_indices = similarities.argsort()[::-1][:top_k]
    return [
        (idx, similarities[idx]) for idx in top_indices if similarities[idx] > 0
    ]

comparison_queries = [
    "neural network training",
    "text processing language",
    "learning from examples",
    "image recognition deep",
]

comparison_results = []
for query in comparison_queries:
    bm25_results = bm25_index.search(query, top_k=5)
    tfidf_results = tfidf_search(query, top_k=5)

    comparison_results.append(
        {"query": query, "bm25": bm25_results, "tfidf": tfidf_results}
    )

# Get all scores for a specific query
query = "neural network learning"

bm25_all_scores = []
tfidf_all_scores = []

query_terms = bm25_index.tokenize(query)
query_vec = tfidf_vectorizer.transform([query])
tfidf_sims = cosine_similarity(query_vec, tfidf_matrix)[0]

for doc_id in range(len(test_corpus)):
    bm25_score_val = bm25_index.score_document(doc_id, query_terms)
    bm25_all_scores.append(bm25_score_val)
    tfidf_all_scores.append(tfidf_sims[doc_id])