Spaces:
Sleeping
Sleeping
| 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 | |