Spaces:
Sleeping
Sleeping
File size: 2,910 Bytes
2290123 bd53034 2290123 bd53034 2290123 bd53034 2290123 bd53034 2290123 | 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 | import math
from typing import Any, Dict, List, Optional
from backend.utils import get_keywords
class BM25Engine:
def __init__(
self, documents: List[str], doc_ids: List[str], metadatas: List[Dict[str, Any]]
):
self.doc_ids = doc_ids
self.documents = documents
self.metadatas = metadatas
self.corpus_size = len(documents)
self.k1 = 1.5
self.b = 0.75
self.tokenized_docs = [get_keywords(doc) for doc in documents]
self.doc_length = [len(doc) for doc in self.tokenized_docs]
self.avg_doc_len = sum(self.doc_length) / max(1, self.corpus_size)
self.doc_tfs = []
self.dfs = {}
self.build_frequencies()
def build_frequencies(self):
for doc in self.tokenized_docs:
tfs = {}
for term in doc:
tfs[term] = tfs.get(term, 0) + 1
self.doc_tfs.append(tfs)
for term in doc:
self.dfs[term] = self.dfs.get(term, 0) + 1
def idf(self, term):
df = self.dfs.get(term, 0)
return math.log(1 + (self.corpus_size - df + 0.5) / (df + 0.5))
def score(self, query: str, doc_index: int) -> float:
query_terms = get_keywords(query)
total_score = 0.0
doc_len = self.doc_length[doc_index]
tfs = self.doc_tfs[doc_index]
for term in query_terms:
tf = tfs.get(term, 0)
if tf > 0:
term_idf = self.idf(term)
# BM25 term weighting formula
numerator = tf * (self.k1 + 1)
denominator = tf + self.k1 * (
1 - self.b + self.b * (doc_len / self.avg_doc_len)
)
total_score += term_idf * (numerator / denominator)
return total_score
def search(
self, query: str, top_k: int = 3, where: Optional[Dict[str, Any]] = None
) -> list:
scores = []
# Score all docs
for idx in range(self.corpus_size):
if where:
match = True
doc_meta = self.metadatas[idx] or {}
for key, val in where.items():
if doc_meta.get(key) != val:
match = False
break
if not match:
continue
s = self.score(query, idx)
scores.append((s, idx))
# Sort descending by score
scores.sort(key=lambda x: x[0], reverse=True)
# Return documents with structured metadata
results = []
for score, idx in scores[:top_k]:
results.append(
{
"id": self.doc_ids[idx],
"text": self.documents[idx],
"metadata": self.metadatas[idx],
"score": score,
}
)
return results
|