Update search_utils.py

search_utils.py

@@ -55,18 +55,22 @@ class MetadataManager:
             raise
 
     def get_metadata(self, global_indices):
-        """Retrieve metadata for given indices"""
+        """Retrieve metadata for given indices with deduplication by title"""
         if isinstance(global_indices, np.ndarray) and global_indices.size == 0:
             return pd.DataFrame(columns=["title", "summary", 'authors', "similarity", "source"])
         
         try:
             # Directly index the DataFrame
             results = self.df.iloc[global_indices].copy()
-            return results.reset_index(drop=True)
+            
+            # Deduplicate by title to avoid near-duplicate results
+            if len(results) > 1:
+                results = results.drop_duplicates(subset=["title"])
+            
+            return results
         except Exception as e:
             logger.error(f"Metadata retrieval failed: {str(e)}")
             return pd.DataFrame(columns=["title", "summary", "similarity", "source", 'authors'])
-    
 
 
 class SemanticSearch:
@@ -134,11 +138,16 @@ class SemanticSearch:
         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."""
+        """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)
     
     def search(self, query, top_k=5):
-        """Search for a query using parallel FAISS shard search."""
+        """Search for a query using parallel FAISS shard search with normalized vectors for proper cosine similarity."""
         self.logger.info(f"Searching for query: '{query}' (top_k={top_k})")
         start_time = time.time()
         if not query:
@@ -150,6 +159,8 @@ class SemanticSearch:
         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)}")
@@ -168,13 +179,29 @@ class SemanticSearch:
                     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"])
+            
         self.logger.info(f"Search found {len(all_global_indices)} results across all shards")
-        results = self._process_results(np.array(all_distances), np.array(all_global_indices), top_k)
+        
+        # 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])
+        
+        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
     
     def _search_shard(self, shard_idx, index, query_embedding, top_k):
-        """Search a single FAISS shard for the query embedding."""
+        """Search a single FAISS shard for the query embedding with proper error handling."""
         if index.ntotal == 0:
             self.logger.warning(f"Skipping empty shard {shard_idx}")
             return None
@@ -182,23 +209,36 @@ class SemanticSearch:
         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_distances = distances[0]
             
+            # Filter out invalid indices (-1 is returned by FAISS for insufficient results)
+            valid_mask = (indices[0] >= 0) & (indices[0] < index.ntotal)
+            valid_indices = indices[0][valid_mask]
+            valid_distances = distances[0][valid_mask]
+            
+            if len(valid_indices) == 0:
+                self.logger.debug(f"Shard {shard_idx}: No valid results found")
+                return None
+                
             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]
-            logger.info(f"Global Indice {global_indices}")
+            
+            # Filter out any invalid global indices (could happen if _global_index validation fails)
+            valid_global = [(d, i) for d, i in zip(valid_distances, global_indices) if i >= 0]
+            if not valid_global:
+                return None
+                
+            final_distances, final_indices = zip(*valid_global)
+            
             self.logger.debug(f"Shard {shard_idx} search completed in {time.time() - shard_start:.3f}s")
-            return valid_distances, global_indices
+            return final_distances, final_indices
         except Exception as e:
             self.logger.error(f"Search failed in shard {shard_idx}: {str(e)}")
             return None
 
     def _process_results(self, distances, global_indices, top_k):
-        """Process raw search results with correct similarity calculation."""
+        """Process raw search results with correct similarity calculation for cosine similarity."""
         process_start = time.time()
         if global_indices.size == 0 or distances.size == 0:
             self.logger.warning("No search results to process")
@@ -221,18 +261,10 @@ class SemanticSearch:
                 results = results.iloc[:min_len]
                 distances = distances[:min_len]
 
-            # Calculate similarity (cosine similarity = inner product for normalized embeddings)
-            results['similarity'] = 1 - (distances / 2)
+            # For inner product with normalized vectors, similarity is directly the distance
+            # (FAISS IP search already returns higher scores for more similar items)
+            results['similarity'] = distances
             
-            # Ensure URL lists are properly formatted
-            # results['source'] = results['source'].apply(
-            #     lambda x: [
-            #         url.strip().rstrip(')')  # Clean trailing parentheses and whitespace
-            #         for url in str(x).split(';')  # Split on semicolons
-            #         if url.strip()  # Remove empty strings
-            #     ] if isinstance(x, (str, list)) else []
-            # )
-
             # Deduplicate and sort
             required_columns = ["title", "summary", "authors", "source", "similarity"]
             pre_dedup = len(results)
@@ -251,4 +283,42 @@ class SemanticSearch:
             
         except Exception as e:
             self.logger.error(f"Result processing failed: {str(e)}", exc_info=True)
-            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
+            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 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)
+        
+        if initial_results.empty or len(initial_results) < 2:
+            return initial_results
+            
+        # 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)