Update search_utils.py

search_utils.py

@@ -80,123 +80,124 @@ class SemanticSearch:
         self.index_shards = []
         self.metadata_mgr = MetadataManager()
         self.shard_sizes = []
-        self.cumulative_offsets = None
-        self.total_vectors = 0
+        
+        # Configure search logger
         self.logger = logging.getLogger("SemanticSearch")
         self.logger.info("Initializing SemanticSearch")
-    
+        
     @st.cache_resource
     def load_model(_self):
         return SentenceTransformer('all-MiniLM-L6-v2')
-    
+
     def initialize_system(self):
         self.logger.info("Loading sentence transformer model")
         start_time = time.time()
         self.model = self.load_model()
         self.logger.info(f"Model loaded in {time.time() - start_time:.2f} seconds")
+        
         self.logger.info("Loading FAISS indices")
         self._load_faiss_shards()
-    
+        
     def _load_faiss_shards(self):
-        """Load FAISS shards concurrently and precompute cumulative offsets for global indexing."""
+        """Load all FAISS index shards"""
         self.logger.info(f"Searching for index files in {self.shard_dir}")
+        
         if not self.shard_dir.exists():
             self.logger.error(f"Shard directory not found: {self.shard_dir}")
             return
-        index_files = sorted(self.shard_dir.glob("*.index"))
+            
+        index_files = list(self.shard_dir.glob("*.index"))
         self.logger.info(f"Found {len(index_files)} index files")
-        self.index_shards = []
+        
         self.shard_sizes = []
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            future_to_file = {
-                executor.submit(self._load_single_index, shard_path): shard_path
-                for shard_path in index_files
-            }
-            for future in concurrent.futures.as_completed(future_to_file):
-                shard_path = future_to_file[future]
-                try:
-                    index, size = future.result()
-                    if index is not None:
-                        self.index_shards.append(index)
-                        self.shard_sizes.append(size)
-                        self.logger.info(f"Loaded index {shard_path.name} with {size} vectors")
-                except Exception as e:
-                    self.logger.error(f"Error loading index {shard_path}: {str(e)}")
+        self.index_shards = []
+        
+        for shard_path in sorted(index_files):
+            try:
+                self.logger.info(f"Loading index: {shard_path}")
+                start_time = time.time()
+                
+                # Log file size
+                file_size_mb = os.path.getsize(shard_path) / (1024 * 1024)
+                self.logger.info(f"Index file size: {file_size_mb:.2f} MB")
+                
+                index = faiss.read_index(str(shard_path))
+                self.index_shards.append(index)
+                self.shard_sizes.append(index.ntotal)
+                
+                self.logger.info(f"Loaded index with {index.ntotal} vectors in {time.time() - start_time:.2f} seconds")
+            except Exception as e:
+                self.logger.error(f"Failed to load index {shard_path}: {str(e)}")
+        
         self.total_vectors = sum(self.shard_sizes)
         self.logger.info(f"Total loaded vectors: {self.total_vectors} across {len(self.index_shards)} shards")
-        self.cumulative_offsets = np.cumsum([0] + self.shard_sizes)
-    
-    def _load_single_index(self, shard_path):
-        """Load a single FAISS index shard."""
-        self.logger.info(f"Loading index: {shard_path}")
-        start_time = time.time()
-        file_size_mb = os.path.getsize(shard_path) / (1024 * 1024)
-        self.logger.info(f"Index file size: {file_size_mb:.2f} MB")
-        index = faiss.read_index(str(shard_path))
-        size = index.ntotal
-        self.logger.info(f"Index loaded in {time.time() - start_time:.2f} seconds")
-        return index, size
     
     def _global_index(self, shard_idx, local_idx):
-        """Convert a local index (within a shard) to a global index using precomputed offsets with validation."""
-        if shard_idx < 0 or shard_idx >= len(self.index_shards):
-            self.logger.error(f"Invalid shard index: {shard_idx}")
-            return -1
-        if local_idx < 0 or local_idx >= self.shard_sizes[shard_idx]:
-            self.logger.warning(f"Local index {local_idx} may be out of bounds for shard {shard_idx}")
-        return int(self.cumulative_offsets[shard_idx] + local_idx)
-    
+        """Convert local index to global index"""
+        return sum(self.shard_sizes[:shard_idx]) + local_idx
+        
+
     def search(self, query, top_k=5):
-        """Search for a query using parallel FAISS shard search with normalized vectors for proper cosine similarity."""
+        """Search with validation"""
         self.logger.info(f"Searching for query: '{query}' (top_k={top_k})")
         start_time = time.time()
+        
         if not query:
             self.logger.warning("Empty query provided")
             return pd.DataFrame()
+            
         if not self.index_shards:
             self.logger.error("No index shards loaded")
             return pd.DataFrame()
+        
         try:
             self.logger.info("Encoding query")
             query_embedding = self.model.encode([query], convert_to_numpy=True)
-            # Normalize query embedding for proper cosine similarity comparison
-            query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
             self.logger.debug(f"Query encoded to shape {query_embedding.shape}")
         except Exception as e:
             self.logger.error(f"Query encoding failed: {str(e)}")
             return pd.DataFrame()
-        
+    
         all_distances = []
         all_global_indices = []
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            futures = {
-                executor.submit(self._search_shard, shard_idx, index, query_embedding, top_k): shard_idx
-                for shard_idx, index in enumerate(self.index_shards)
-            }
-            for future in concurrent.futures.as_completed(futures):
-                result = future.result()
-                if result is not None:
-                    distances_part, global_indices_part = result
-                    all_distances.extend(distances_part)
-                    all_global_indices.extend(global_indices_part)
-        
-        # If no results found across all shards
-        if not all_global_indices:
-            self.logger.warning("No results found across any shards")
-            return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
-            
+    
+        # Search with index validation
+        self.logger.info(f"Searching across {len(self.index_shards)} shards")
+        for shard_idx, index in enumerate(self.index_shards):
+            if index.ntotal == 0:
+                self.logger.warning(f"Skipping empty shard {shard_idx}")
+                continue
+                
+            try:
+                shard_start = time.time()
+                distances, indices = index.search(query_embedding, top_k)
+                
+                valid_mask = (indices[0] >= 0) & (indices[0] < index.ntotal)
+                valid_indices = indices[0][valid_mask].tolist()
+                valid_distances = distances[0][valid_mask].tolist()
+                
+                if len(valid_indices) != top_k:
+                    self.logger.debug(f"Shard {shard_idx}: Found {len(valid_indices)} valid results out of {top_k}")
+                
+                global_indices = [self._global_index(shard_idx, idx) for idx in valid_indices]
+                
+                all_distances.extend(valid_distances)
+                all_global_indices.extend(global_indices)
+                
+                self.logger.debug(f"Shard {shard_idx} search completed in {time.time() - shard_start:.3f}s")
+            except Exception as e:
+                self.logger.error(f"Search failed in shard {shard_idx}: {str(e)}")
+                continue
+    
         self.logger.info(f"Search found {len(all_global_indices)} results across all shards")
         
-        # Sort all results by distance before processing
-        combined = list(zip(all_distances, all_global_indices))
-        combined.sort(reverse=True)  # Sort by distance (higher is better for cosine similarity)
-        sorted_distances, sorted_indices = zip(*combined)
-        
-        # Limit to top-k across all shards
-        top_distances = np.array(sorted_distances[:top_k])
-        top_indices = np.array(sorted_indices[:top_k])
+        # Process results
+        results = self._process_results(
+            np.array(all_distances), 
+            np.array(all_global_indices), 
+            top_k
+        )
         
-        results = self._process_results(top_distances, top_indices, top_k)
         self.logger.info(f"Search completed in {time.time() - start_time:.2f} seconds with {len(results)} final results")
         return results
     
@@ -285,40 +286,77 @@ class SemanticSearch:
             self.logger.error(f"Result processing failed: {str(e)}", exc_info=True)
             return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
     
-    def search_with_threshold(self, query, top_k=10, similarity_threshold=0.6):
-        """
-        Search with a fixed similarity threshold, returning only results above the threshold.
-        For cosine similarity with normalized vectors, threshold should be between 0 and 1.
-        """
-        # Get more results initially to ensure we have enough after filtering
-        initial_results = self.search(query, top_k=top_k*2)
-        
-        if initial_results.empty:
-            return initial_results
-            
-        # Filter by similarity threshold
-        filtered_results = initial_results[initial_results['similarity'] >= similarity_threshold]
-        
-        # Return top-k of filtered results
-        return filtered_results.head(top_k).reset_index(drop=True)
+
+    
+
+    
+
+    def _process_results(self, distances, global_indices, top_k):
+        """Process raw search results into formatted DataFrame"""
+        process_start = time.time()
         
-    def search_with_adaptive_threshold(self, query, top_k=10, percentile=75):
-        """
-        Search with an adaptive threshold based on the distribution of similarity scores.
-        Returns results above the specified percentile of similarity scores.
-        """
-        # Get more results initially to determine distribution
-        initial_results = self.search(query, top_k=top_k*3)
+        # Proper numpy array emptiness checks
+        if global_indices.size == 0 or distances.size == 0:
+            self.logger.warning("No search results to process")
+            return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
         
-        if initial_results.empty or len(initial_results) < 2:
-            return initial_results
+        try:
+            # Get metadata for matched indices
+            self.logger.info(f"Retrieving metadata for {len(global_indices)} indices")
+            metadata_start = time.time()
+            results = self.metadata_mgr.get_metadata(global_indices)
+            self.logger.info(f"Metadata retrieved in {time.time() - metadata_start:.2f}s, got {len(results)} records")
             
-        # Calculate threshold based on percentile of similarity scores
-        threshold = np.percentile(initial_results['similarity'], percentile)
-        self.logger.info(f"Adaptive threshold set to {threshold:.4f} (percentile: {percentile})")
-        
-        # Filter results above threshold
-        filtered_results = initial_results[initial_results['similarity'] >= threshold]
-        
-        # Return top-k of filtered results
-        return filtered_results.head(top_k).reset_index(drop=True)
+            # Empty results check
+            if len(results) == 0:
+                self.logger.warning("No metadata found for indices")
+                return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
+                
+            # Ensure distances match results length
+            if len(results) != len(distances):
+                self.logger.warning(f"Mismatch between distances ({len(distances)}) and results ({len(results)})")
+                
+                if len(results) < len(distances):
+                    self.logger.info("Truncating distances array to match results length")
+                    distances = distances[:len(results)]
+                else:
+                    # Should not happen but handle it anyway
+                    self.logger.error("More results than distances - this shouldn't happen")
+                    distances = np.pad(distances, (0, len(results) - len(distances)), 'constant', constant_values=1.0)
+                
+            # Calculate similarity scores
+            self.logger.debug("Calculating similarity scores")
+            results['similarity'] = 1 - (distances / 2)
+            
+            # Log similarity statistics
+            if not results.empty:
+                self.logger.debug(f"Similarity stats: min={results['similarity'].min():.3f}, " +
+                                 f"max={results['similarity'].max():.3f}, " +
+                                 f"mean={results['similarity'].mean():.3f}")
+
+            
+            # Deduplicate and sort results
+            pre_dedup = len(results)
+            results = results.drop_duplicates(subset=["title", "source"]).sort_values("similarity", ascending=False).head(top_k)
+            post_dedup = len(results)
+            
+            if pre_dedup > post_dedup:
+                self.logger.info(f"Removed {pre_dedup - post_dedup} duplicate results")
+            
+            self.logger.info(f"Results processed in {time.time() - process_start:.2f}s, returning {len(results)} items")
+            return results.reset_index(drop=True)
+
+            # Add URL resolution for final results only
+            final_results = results.sort_values("similarity", ascending=False).head(top_k)
+            
+            # Resolve URLs for top results only
+            # final_results['source'] =
+    
+            # Deduplicate based on title only
+            final_results = final_results.drop_duplicates(subset=["title"]).head(top_k)
+            
+            return final_results.reset_index(drop=True)
+            
+        except Exception as e:
+            self.logger.error(f"Result processing failed: {str(e)}", exc_info=True)
+            return pd.DataFrame(columns=["title", "summary", "similarity"])