Spaces:
Runtime error
Runtime error
| """ | |
| BM25 keyword-based search for hybrid retrieval. | |
| Provides exact keyword matching to complement semantic search. | |
| """ | |
| from typing import List, Dict | |
| from rank_bm25 import BM25Okapi | |
| import numpy as np | |
| class BM25Search: | |
| """BM25 keyword-based search engine.""" | |
| def __init__(self): | |
| """Initialize BM25 search.""" | |
| self.corpus = [] | |
| self.tokenized_corpus = [] | |
| self.bm25 = None | |
| self.chunks = [] | |
| print("[BM25] Initialized") | |
| def index_documents(self, chunks: List[Dict]): | |
| """ | |
| Index documents for BM25 search. | |
| Args: | |
| chunks: List of chunk dictionaries with 'content' key | |
| """ | |
| self.chunks = chunks | |
| self.corpus = [chunk['content'] for chunk in chunks] | |
| # Simple tokenization (lowercase + split) | |
| self.tokenized_corpus = [ | |
| doc.lower().split() for doc in self.corpus | |
| ] | |
| # Build BM25 index | |
| if self.tokenized_corpus: | |
| self.bm25 = BM25Okapi(self.tokenized_corpus) | |
| print(f"[BM25] Indexed {len(self.chunks)} documents") | |
| def search( | |
| self, | |
| query: str, | |
| top_k: int = 30 | |
| ) -> List[Dict]: | |
| """ | |
| Search documents using BM25. | |
| Args: | |
| query: Search query | |
| top_k: Number of results to return | |
| Returns: | |
| List of chunks sorted by BM25 score | |
| """ | |
| if not self.bm25: | |
| return [] | |
| # Tokenize query | |
| tokenized_query = query.lower().split() | |
| # Get BM25 scores | |
| scores = self.bm25.get_scores(tokenized_query) | |
| # Get top_k indices | |
| top_indices = np.argsort(scores)[::-1][:top_k] | |
| # Return chunks with scores | |
| results = [] | |
| for idx in top_indices: | |
| if scores[idx] > 0: # Only return non-zero scores | |
| chunk = self.chunks[idx].copy() | |
| chunk['bm25_score'] = float(scores[idx]) | |
| results.append(chunk) | |
| return results | |
| def get_score(self, query: str, document: str) -> float: | |
| """ | |
| Get BM25 score for a single query-document pair. | |
| Args: | |
| query: Search query | |
| document: Document text | |
| Returns: | |
| BM25 relevance score | |
| """ | |
| tokenized_query = query.lower().split() | |
| tokenized_doc = document.lower().split() | |
| # Create temporary BM25 for single document | |
| temp_bm25 = BM25Okapi([tokenized_doc]) | |
| score = temp_bm25.get_scores(tokenized_query)[0] | |
| return float(score) | |
| class HybridSearch: | |
| """Combines semantic and keyword search.""" | |
| def __init__( | |
| self, | |
| semantic_weight: float = 0.7, | |
| keyword_weight: float = 0.3 | |
| ): | |
| """ | |
| Initialize hybrid search. | |
| Args: | |
| semantic_weight: Weight for semantic search (0-1) | |
| keyword_weight: Weight for keyword search (0-1) | |
| """ | |
| self.semantic_weight = semantic_weight | |
| self.keyword_weight = keyword_weight | |
| self.bm25 = BM25Search() | |
| print(f"[Hybrid Search] Initialized (semantic: {semantic_weight}, keyword: {keyword_weight})") | |
| def combine_results( | |
| self, | |
| semantic_results: List[Dict], | |
| keyword_results: List[Dict], | |
| top_k: int = 10 | |
| ) -> List[Dict]: | |
| """ | |
| Combine and rerank results from semantic and keyword search. | |
| Args: | |
| semantic_results: Results from semantic search (with 'distance' scores) | |
| keyword_results: Results from BM25 search (with 'bm25_score') | |
| top_k: Number of final results | |
| Returns: | |
| Combined and reranked results | |
| """ | |
| # Normalize scores to 0-1 range | |
| def normalize_scores(results, score_key): | |
| if not results: | |
| return results | |
| scores = [r.get(score_key, 0) for r in results] | |
| min_score = min(scores) | |
| max_score = max(scores) | |
| if max_score == min_score: | |
| return results | |
| for r in results: | |
| r[f'{score_key}_normalized'] = ( | |
| (r.get(score_key, 0) - min_score) / (max_score - min_score) | |
| ) | |
| return results | |
| # For semantic search, lower distance = higher relevance | |
| # Need to invert: score = 1 - normalized_distance | |
| for r in semantic_results: | |
| if 'distance' in r: | |
| r['semantic_score'] = r['distance'] # Will normalize below | |
| semantic_results = normalize_scores(semantic_results, 'semantic_score') | |
| keyword_results = normalize_scores(keyword_results, 'bm25_score') | |
| # Invert semantic scores (lower distance = better) | |
| for r in semantic_results: | |
| if 'semantic_score_normalized' in r: | |
| r['semantic_score_normalized'] = 1 - r['semantic_score_normalized'] | |
| # Merge results by chunk ID or content | |
| merged = {} | |
| for chunk in semantic_results: | |
| chunk_id = chunk.get('metadata', {}).get('document_id', '') + '_' + str(chunk.get('metadata', {}).get('chunk_index', '')) | |
| merged[chunk_id] = chunk.copy() | |
| merged[chunk_id]['hybrid_score'] = ( | |
| self.semantic_weight * chunk.get('semantic_score_normalized', 0) | |
| ) | |
| for chunk in keyword_results: | |
| chunk_id = chunk.get('metadata', {}).get('document_id', '') + '_' + str(chunk.get('metadata', {}).get('chunk_index', '')) | |
| if chunk_id in merged: | |
| merged[chunk_id]['hybrid_score'] += ( | |
| self.keyword_weight * chunk.get('bm25_score_normalized', 0) | |
| ) | |
| else: | |
| merged[chunk_id] = chunk.copy() | |
| merged[chunk_id]['hybrid_score'] = ( | |
| self.keyword_weight * chunk.get('bm25_score_normalized', 0) | |
| ) | |
| # Sort by hybrid score | |
| results = sorted( | |
| merged.values(), | |
| key=lambda x: x.get('hybrid_score', 0), | |
| reverse=True | |
| ) | |
| return results[:top_k] | |
| # Global hybrid search instance | |
| hybrid_search = HybridSearch() | |