Spaces:
Running
Running
| """Sparse keyword search using BM25 (Best Matching 25). | |
| Implemented cleanly in pure Python with optional rank_bm25 fallback: | |
| Zero external dependencies required! | |
| Why BM25? | |
| While vector embeddings capture semantic meaning (synonyms, conceptual similarity), | |
| BM25 captures exact lexical matches (rare keywords, acronyms, specific numbers). | |
| """ | |
| import math | |
| import re | |
| from collections import Counter | |
| TOKEN_RE = re.compile(r"[A-Za-z0-9_\-]+") | |
| def tokenize(text): | |
| """Clean and tokenize text into lowercase words/terms.""" | |
| if not text: | |
| return [] | |
| return [t.lower() for t in TOKEN_RE.findall(text)] | |
| class SimpleBM25: | |
| """Pure Python implementation of Okapi BM25. | |
| Formula: | |
| IDF(q) = ln((N - n(q) + 0.5) / (n(q) + 0.5) + 1) | |
| Score(D, Q) = sum( IDF(q) * (f(q, D) * (k1 + 1)) / (f(q, D) + k1 * (1 - b + b * (|D| / avgdl))) ) | |
| """ | |
| def __init__(self, corpus, k1=1.5, b=0.75): | |
| self.k1 = k1 | |
| self.b = b | |
| self.corpus_size = len(corpus) | |
| self.doc_lengths = [len(doc) for doc in corpus] | |
| self.avg_doc_len = sum(self.doc_lengths) / max(1, self.corpus_size) | |
| # Document frequencies for each term | |
| self.doc_freqs = Counter() | |
| self.doc_term_counts = [] | |
| for doc in corpus: | |
| counts = Counter(doc) | |
| self.doc_term_counts.append(counts) | |
| for term in counts.keys(): | |
| self.doc_freqs[term] += 1 | |
| def get_scores(self, query_tokens): | |
| """Calculate BM25 relevance score for each document in the corpus.""" | |
| scores = [0.0] * self.corpus_size | |
| for token in query_tokens: | |
| if token not in self.doc_freqs: | |
| continue | |
| # Inverse Document Frequency (IDF) | |
| df = self.doc_freqs[token] | |
| idf = math.log((self.corpus_size - df + 0.5) / (df + 0.5) + 1.0) | |
| # Term frequency & length normalization for each doc | |
| for doc_idx, term_counts in enumerate(self.doc_term_counts): | |
| tf = term_counts.get(token, 0) | |
| if tf == 0: | |
| continue | |
| doc_len = self.doc_lengths[doc_idx] | |
| numerator = tf * (self.k1 + 1.0) | |
| denominator = tf + self.k1 * (1.0 - self.b + self.b * (doc_len / max(1.0, self.avg_doc_len))) | |
| scores[doc_idx] += idf * (numerator / denominator) | |
| return scores | |
| class BM25SearchIndex: | |
| """Keyword search index over chunks.""" | |
| def __init__(self): | |
| self.chunks = [] | |
| self.corpus_tokens = [] | |
| self.bm25 = None | |
| def build(self, chunks): | |
| """Index a list of TextChunk objects.""" | |
| if not chunks: | |
| self.chunks = [] | |
| self.corpus_tokens = [] | |
| self.bm25 = None | |
| return | |
| self.chunks = chunks | |
| self.corpus_tokens = [] | |
| for chunk in chunks: | |
| # Include text and document title for matching | |
| combined_text = f"{chunk.title} {chunk.text}" | |
| self.corpus_tokens.append(tokenize(combined_text)) | |
| try: | |
| from rank_bm25 import BM25Okapi | |
| self.bm25 = BM25Okapi(self.corpus_tokens) | |
| except ImportError: | |
| # Use pure Python BM25 fallback | |
| self.bm25 = SimpleBM25(self.corpus_tokens) | |
| def search(self, query, top_k=5): | |
| """Find the top_k best matching chunks by BM25 keyword relevance. | |
| Returns: | |
| List of dicts: [{"chunk": TextChunk, "score": float, "source": "bm25"}] | |
| """ | |
| if not self.bm25 or not self.chunks: | |
| return [] | |
| query_tokens = tokenize(query) | |
| if not query_tokens: | |
| return [] | |
| scores = self.bm25.get_scores(query_tokens) | |
| scored = list(enumerate(scores)) | |
| scored.sort(key=lambda x: x[1], reverse=True) | |
| results = [] | |
| for idx, score in scored[:top_k]: | |
| if score <= 0: | |
| continue | |
| results.append({ | |
| "chunk": self.chunks[idx], | |
| "score": round(float(score), 4), | |
| "source": "bm25", | |
| }) | |
| return results | |
| def __len__(self): | |
| return len(self.chunks) | |