Spaces:
Running
Running
| """ | |
| retriever.py — Production Hybrid C++17 AVX2 SIMD & Okapi BM25 RAG Retriever with RRF & Reranking. | |
| Combines: | |
| 1. Dense vector similarity via C++17 AVX2 SIMD dot-product engine. | |
| 2. Sparse keyword matching via Okapi BM25 engine. | |
| 3. Reciprocal Rank Fusion (RRF, k=60) for balanced multi-modal candidate selection. | |
| 4. Cross-Attention neural reranker for precision context distillation. | |
| """ | |
| import numpy as np | |
| import time | |
| import logging | |
| from typing import List, Optional, Tuple, Dict | |
| from dataclasses import dataclass, field | |
| from ..config import RAG_TOP_K, RAG_SIMILARITY_THRESHOLD, EMBEDDING_DIM | |
| from .pdf_chunker import DocumentChunk | |
| from .bm25_retriever import BM25Retriever, BM25Match | |
| from .reranker import CrossAttentionReranker, RerankedChunk | |
| logger = logging.getLogger(__name__) | |
| # Try importing the compiled SIMD engine; fall back to NumPy | |
| try: | |
| import simd_engine | |
| SIMD_AVAILABLE = True | |
| logger.info("C++17 AVX2 SIMD engine loaded successfully") | |
| except ImportError: | |
| SIMD_AVAILABLE = False | |
| logger.warning("SIMD engine not available, using NumPy fallback") | |
| class RetrievalResult: | |
| """Result of a retrieval query.""" | |
| chunks: List[DocumentChunk] | |
| scores: List[float] | |
| query_text: str = "" | |
| search_latency_ms: float = 0.0 | |
| total_chunks_searched: int = 0 | |
| used_simd: bool = False | |
| used_hybrid: bool = False | |
| used_reranker: bool = False | |
| reranked_details: List[RerankedChunk] = field(default_factory=list) | |
| class SIMDRetriever: | |
| """ | |
| Production Hybrid Retriever fusing C++17 AVX2 SIMD dense search, | |
| Okapi BM25 sparse keyword search, RRF fusion, and two-stage cross-attention reranking. | |
| """ | |
| def __init__( | |
| self, | |
| top_k: int = RAG_TOP_K, | |
| similarity_threshold: float = RAG_SIMILARITY_THRESHOLD, | |
| ): | |
| self.top_k = top_k | |
| self.similarity_threshold = similarity_threshold | |
| # Chunk store | |
| self._chunks: List[DocumentChunk] = [] | |
| self._embeddings: Optional[np.ndarray] = None # shape: (n, dim) | |
| # Hybrid components | |
| self._bm25 = BM25Retriever() | |
| self._reranker = CrossAttentionReranker() | |
| self._total_queries = 0 | |
| self._total_latency_ms = 0.0 | |
| def ingest(self, chunks: List[DocumentChunk], embeddings: np.ndarray): | |
| """ | |
| Ingest pre-embedded chunks into both dense AVX2 store and sparse BM25 index. | |
| Zero-double-embedding principle: each chunk embedded exactly once at ingestion. | |
| """ | |
| if len(chunks) != embeddings.shape[0]: | |
| raise ValueError( | |
| f"Chunks count ({len(chunks)}) != embeddings count ({embeddings.shape[0]})" | |
| ) | |
| embeddings = np.ascontiguousarray(embeddings, dtype=np.float32) | |
| # L2 normalize | |
| norms = np.linalg.norm(embeddings, axis=1, keepdims=True) | |
| norms = np.maximum(norms, 1e-9) | |
| embeddings = embeddings / norms | |
| self._chunks = list(chunks) | |
| self._embeddings = embeddings | |
| # Build BM25 sparse keyword index over the chunks | |
| self._bm25.build_index(chunks) | |
| logger.info( | |
| f"Successfully ingested {len(chunks)} chunks into C++17 AVX2 vector store " | |
| f"and BM25 keyword index." | |
| ) | |
| def search( | |
| self, | |
| query_embedding: np.ndarray, | |
| top_k: Optional[int] = None, | |
| query_text: Optional[str] = None, | |
| ) -> RetrievalResult: | |
| """ | |
| Search for the top-k most relevant chunks. | |
| If query_text is provided, executes full Hybrid Search (BM25 + AVX2 SIMD + RRF + Reranking). | |
| If query_text is None, executes pure AVX2 SIMD dense search for backward compatibility. | |
| """ | |
| if self._embeddings is None or len(self._chunks) == 0: | |
| return RetrievalResult( | |
| chunks=[], scores=[], query_text=query_text or "", | |
| total_chunks_searched=0, used_simd=False, | |
| ) | |
| k = top_k or self.top_k | |
| query_vec = np.ascontiguousarray(query_embedding.flatten(), dtype=np.float32) | |
| start = time.perf_counter() | |
| # 1. Dense SIMD / NumPy candidate search | |
| dense_candidates_count = min(len(self._chunks), max(k * 3, 15)) | |
| if SIMD_AVAILABLE: | |
| result = simd_engine.top_k_search( | |
| query_vec, | |
| self._embeddings, | |
| dense_candidates_count, | |
| -1.0, | |
| ) | |
| dense_indices = list(result.indices) | |
| dense_scores = list(result.scores) | |
| used_simd = True | |
| else: | |
| similarities = self._embeddings @ query_vec | |
| top_idx = np.argsort(similarities)[::-1][:dense_candidates_count] | |
| dense_indices = top_idx.tolist() | |
| dense_scores = similarities[top_idx].tolist() | |
| used_simd = False | |
| dense_scores_map = {idx: float(score) for idx, score in zip(dense_indices, dense_scores)} | |
| # 2. If query_text is provided, perform Hybrid BM25 Fusion + Cross-Attention Reranking | |
| if query_text and query_text.strip(): | |
| # Sparse BM25 Search | |
| tokens = BM25Retriever.tokenize(query_text) | |
| sparse_matches = self._bm25.search(tokens, top_k=dense_candidates_count) | |
| # Reciprocal Rank Fusion (RRF) with standard constant 60 | |
| rrf_constant = 60 | |
| rrf_scores: Dict[int, float] = {} | |
| # Add Dense ranks (weight 0.60) | |
| for rank, idx in enumerate(dense_indices): | |
| rrf_scores[idx] = rrf_scores.get(idx, 0.0) + (0.60 / (rrf_constant + rank + 1)) | |
| # Add Sparse BM25 ranks (weight 0.40) | |
| for rank, bm25_match in enumerate(sparse_matches): | |
| idx = bm25_match.chunk_index | |
| rrf_scores[idx] = rrf_scores.get(idx, 0.0) + (0.40 / (rrf_constant + rank + 1)) | |
| # Sort fused candidates by RRF score | |
| sorted_candidate_ids = sorted( | |
| rrf_scores.keys(), key=lambda doc_id: rrf_scores[doc_id], reverse=True | |
| )[:dense_candidates_count] | |
| candidate_pairs = [ | |
| (self._chunks[idx], dense_scores_map.get(idx, 0.35)) | |
| for idx in sorted_candidate_ids | |
| ] | |
| # 3. Two-Stage Cross-Attention Reranking | |
| reranked = self._reranker.rerank(query=query_text, candidate_chunks=candidate_pairs, top_k=k) | |
| elapsed_ms = (time.perf_counter() - start) * 1000.0 | |
| self._total_queries += 1 | |
| self._total_latency_ms += elapsed_ms | |
| result_chunks = [r.chunk for r in reranked] | |
| result_scores = [r.rerank_score for r in reranked] | |
| return RetrievalResult( | |
| chunks=result_chunks, | |
| scores=result_scores, | |
| query_text=query_text, | |
| search_latency_ms=elapsed_ms, | |
| total_chunks_searched=len(self._chunks), | |
| used_simd=used_simd, | |
| used_hybrid=True, | |
| used_reranker=True, | |
| reranked_details=reranked, | |
| ) | |
| # Fallback: Pure Dense Vector Search (e.g. for vector-only tests) | |
| elapsed_ms = (time.perf_counter() - start) * 1000.0 | |
| self._total_queries += 1 | |
| self._total_latency_ms += elapsed_ms | |
| result_chunks = [self._chunks[i] for i in dense_indices[:k]] | |
| result_scores = [float(s) for s in dense_scores[:k]] | |
| return RetrievalResult( | |
| chunks=result_chunks, | |
| scores=result_scores, | |
| query_text="", | |
| search_latency_ms=elapsed_ms, | |
| total_chunks_searched=len(self._chunks), | |
| used_simd=used_simd, | |
| used_hybrid=False, | |
| used_reranker=False, | |
| ) | |
| def hybrid_search( | |
| self, | |
| query_embedding: np.ndarray, | |
| query_text: str, | |
| top_k: Optional[int] = None, | |
| ) -> RetrievalResult: | |
| """Explicit entrypoint for full Hybrid (BM25 + AVX2 + RRF + Reranker) search.""" | |
| return self.search( | |
| query_embedding=query_embedding, | |
| top_k=top_k, | |
| query_text=query_text, | |
| ) | |
| def retrieve(self, query_embedding: np.ndarray, top_k: int = 5) -> RetrievalResult: | |
| """Legacy helper for backward compatibility.""" | |
| return self.search(query_embedding=query_embedding, top_k=top_k) | |
| def chunk_count(self) -> int: | |
| return len(self._chunks) | |
| def stats(self) -> dict: | |
| avg_latency = self._total_latency_ms / max(self._total_queries, 1) | |
| return { | |
| "total_queries": self._total_queries, | |
| "avg_search_latency_ms": round(avg_latency, 4), | |
| "total_chunks": len(self._chunks), | |
| "simd_available": SIMD_AVAILABLE, | |
| "hybrid_bm25_ready": True, | |
| "reranker_ready": True, | |
| } | |