Implementando busca por palavras-chave com BM25

# Install: pip install rank-bm25
from rank_bm25 import BM25Okapi

# BM25Okapi is the most common variant
# BM25L and BM25Plus handle very short documents better
# For most RAG use cases BM25Okapi is the right choice

corpus = [
    'Python decorator pattern explained with examples',
    'How to use context managers in Python',
    'JavaScript async await tutorial',
]
tokenized = [doc.lower().split() for doc in corpus]
bm25 = BM25Okapi(tokenized)
print('Index built with', len(corpus), 'documents')

import re
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer

STOP_WORDS = set(stopwords.words('english'))
stemmer = PorterStemmer()

def tokenize(text: str) -> list[str]:
    text = text.lower()
    text = re.sub(r'[^a-z0-9\s]', ' ', text)
    tokens = text.split()
    tokens = [t for t in tokens if t not in STOP_WORDS and len(t) > 1]
    tokens = [stemmer.stem(t) for t in tokens]
    return tokens

print(tokenize('Running Python decorators efficiently in production!'))
# ['run', 'python', 'decor', 'effici', 'product']

from rank_bm25 import BM25Okapi

def build_bm25_index(documents: list[str]):
    tokenized = [tokenize(doc) for doc in documents]
    bm25 = BM25Okapi(tokenized)
    return bm25, tokenized

# Example with a small corpus
docs = [
    'Vector databases store dense embeddings for similarity search',
    'BM25 is a sparse keyword retrieval algorithm used in search engines',
    'Hybrid search combines dense and sparse retrieval for better recall',
    'PostgreSQL supports vector search via the pgvector extension',
]
bm25, tokenized = build_bm25_index(docs)
print(f'Index contains {bm25.corpus_size} documents')

def bm25_search(bm25, documents: list[str], query: str, top_k: int = 3):
    query_tokens = tokenize(query)
    scores = bm25.get_scores(query_tokens)

    # Get indices sorted by score descending
    ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)

    results = []
    for idx, score in ranked[:top_k]:
        results.append({
            'document': documents[idx],
            'score': round(score, 4),
            'rank': len(results) + 1,
        })
    return results

results = bm25_search(bm25, docs, 'sparse keyword search engine')
for r in results:
    print(f"Rank {r['rank']} (score {r['score']}): {r['document'][:60]}")

from rank_bm25 import BM25Okapi

# Default hyperparameters — good starting point
bm25_default = BM25Okapi(tokenized, k1=1.5, b=0.75)

# Tuned for short document chunks
bm25_short = BM25Okapi(tokenized, k1=1.2, b=0.3)

# Tuned for long documents
bm25_long = BM25Okapi(tokenized, k1=2.0, b=0.9)

# Always benchmark hyperparameters against a golden eval set
# before deploying to production

import re

def technical_tokenize(text: str) -> list[str]:
    text = text.lower()
    # preserve underscores in snake_case and dots in version numbers
    text = re.sub(r'[^a-z0-9_.\s]', ' ', text)
    tokens = text.split()
    # keep tokens that look like identifiers (contain _ or .)
    tokens = [
        t for t in tokens
        if len(t) > 1 and t not in STOP_WORDS
    ]
    return tokens

print(technical_tokenize('Install pgvector 0.5.1 extension in PostgreSQL 16'))
# ['pgvector', '0.5.1', 'extension', 'postgresql', '16']

import pickle

def save_bm25_index(bm25, documents: list[str], path: str):
    with open(path, 'wb') as f:
        pickle.dump({'bm25': bm25, 'documents': documents}, f)
    print(f'Index saved to {path}')

def load_bm25_index(path: str):
    with open(path, 'rb') as f:
        data = pickle.load(f)
    return data['bm25'], data['documents']

save_bm25_index(bm25, docs, '/tmp/bm25_index.pkl')
bm25_loaded, docs_loaded = load_bm25_index('/tmp/bm25_index.pkl')

import threading

class SwappableBM25Index:
    def __init__(self):
        self._index = None
        self._docs = []
        self._lock = threading.RLock()

    def rebuild(self, new_docs: list[str]):
        tokenized = [tokenize(d) for d in new_docs]
        new_index = BM25Okapi(tokenized)
        with self._lock:
            self._index = new_index
            self._docs = new_docs
        print(f'Index rebuilt with {len(new_docs)} documents')

    def search(self, query: str, top_k: int = 5):
        with self._lock:
            return bm25_search(self._index, self._docs, query, top_k)

from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_core.documents import Document

langchain_docs = [Document(page_content=d) for d in docs]

bm25_retriever = BM25Retriever.from_documents(langchain_docs)
bm25_retriever.k = 5

# Combine with a vector retriever (assuming vector_retriever is already defined)
# ensemble = EnsembleRetriever(
#     retrievers=[bm25_retriever, vector_retriever],
#     weights=[0.4, 0.6],  # 40% BM25, 60% dense
# )

results = bm25_retriever.invoke('sparse keyword search')
for doc in results:
    print(doc.page_content[:80])

def hit_rate_at_k(bm25, documents, queries, relevant_docs, k=5):
    hits = 0
    for query, relevant in zip(queries, relevant_docs):
        results = bm25_search(bm25, documents, query, top_k=k)
        retrieved = [r['document'] for r in results]
        if relevant in retrieved:
            hits += 1
    return hits / len(queries)

# Example evaluation
test_queries = ['BM25 algorithm', 'hybrid search systems']
test_relevant = [
    'BM25 is a sparse keyword retrieval algorithm used in search engines',
    'Hybrid search combines dense and sparse retrieval for better recall',
]
hit_rate = hit_rate_at_k(bm25, docs, test_queries, test_relevant, k=3)
print(f'Hit rate @3: {hit_rate:.2%}')