Spaces:
Running
Running
File size: 4,215 Bytes
e43c82d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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)
|