Create search_utils.py

search_utils.py

@@ -0,0 +1,64 @@
+import numpy as np
+import pandas as pd
+import faiss
+from pathlib import Path
+from sentence_transformers import SentenceTransformer, util
+import streamlit as st
+
+class SemanticSearch:
+    def __init__(self, shard_dir="compressed_shards"):
+        self.shard_dir = Path(shard_dir)
+        self.shard_dir.mkdir(exist_ok=True, parents=True)
+        self.model = None
+        self.index_shards = []
+        
+    @st.cache_resource
+    def load_model(_self):
+        return SentenceTransformer('all-MiniLM-L6-v2')
+        
+    def initialize_system(self):
+        self.model = self.load_model()
+        self._load_index_shards()
+
+    def _load_index_shards(self):
+        """Load FAISS shards directly from local directory"""
+        for shard_path in sorted(self.shard_dir.glob("*.index")):
+            self.index_shards.append(faiss.read_index(str(shard_path)))
+
+    def search(self, query, top_k=5):
+        """Search across all shards"""
+        query_embedding = self.model.encode([query], convert_to_numpy=True)
+        all_scores = []
+        all_indices = []
+        
+        for shard_idx, index in enumerate(self.index_shards):
+            distances, indices = index.search(query_embedding, top_k)
+            # Convert local indices to global shard offsets
+            global_indices = [
+                self._calculate_global_index(shard_idx, idx)
+                for idx in indices[0]
+            ]
+            all_scores.extend(distances[0])
+            all_indices.extend(global_indices)
+            
+        return self._process_results(np.array(all_scores), np.array(all_indices), top_k)
+
+    def _calculate_global_index(self, shard_idx, local_idx):
+        """Convert shard-local index to global index"""
+        # Implement your specific shard indexing logic here
+        # Example: return f"{shard_idx}-{local_idx}"
+        return local_idx  # Simple version if using unique IDs
+
+    def _process_results(self, distances, indices, top_k):
+        """Format search results"""
+        results = pd.DataFrame({
+            'global_index': indices,
+            'similarity': 1 - (distances / 2)  # L2 to cosine approximation
+        })
+        return results.sort_values('similarity', ascending=False).head(top_k)
+
+    def search_with_threshold(self, query, top_k=5, similarity_threshold=0.6):
+        """Threshold-filtered search"""
+        results = self.search(query, top_k*2)
+        filtered = results[results['similarity'] > similarity_threshold].head(top_k)
+        return filtered.reset_index(drop=True)