Update search_utils.py

search_utils.py

@@ -2,9 +2,22 @@ import numpy as np
 import pandas as pd
 import faiss
 import zipfile
+import logging
 from pathlib import Path
 from sentence_transformers import SentenceTransformer, util
 import streamlit as st
+import time
+import os
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger("MetadataManager")
 
 class MetadataManager:
     def __init__(self):
@@ -12,45 +25,94 @@ class MetadataManager:
         self.shard_map = {}
         self.loaded_shards = {}
         self.total_docs = 0
-        self._ensure_unzipped()  # Removed Streamlit elements from here
+        
+        logger.info("Initializing MetadataManager")
+        self._ensure_unzipped()
         self._build_shard_map()
+        logger.info(f"Total documents indexed: {self.total_docs}")
+        logger.info(f"Total shards found: {len(self.shard_map)}")
 
     def _ensure_unzipped(self):
         """Handle ZIP extraction without Streamlit elements"""
+        logger.info(f"Checking for shard directory: {self.shard_dir}")
         if not self.shard_dir.exists():
             zip_path = Path("metadata_shards.zip")
+            logger.info(f"Shard directory not found, looking for zip file: {zip_path}")
             if zip_path.exists():
+                logger.info(f"Extracting from zip file: {zip_path}")
+                start_time = time.time()
                 with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                     zip_ref.extractall(self.shard_dir)
+                logger.info(f"Extraction completed in {time.time() - start_time:.2f} seconds")
             else:
-                raise FileNotFoundError("Metadata ZIP file not found")
+                error_msg = "Metadata ZIP file not found"
+                logger.error(error_msg)
+                raise FileNotFoundError(error_msg)
+        else:
+            logger.info("Shard directory exists, skipping extraction")
                 
     def _build_shard_map(self):
         """Create index range to shard mapping"""
+        logger.info("Building shard map from parquet files")
         self.total_docs = 0
-        for f in sorted(self.shard_dir.glob("*.parquet")):
-            parts = f.stem.split("_")
-            start = int(parts[1])
-            end = int(parts[2])
-            self.shard_map[(start, end)] = f.name
-            self.total_docs = max(self.total_docs, end + 1)
+        shard_files = list(self.shard_dir.glob("*.parquet"))
+        logger.info(f"Found {len(shard_files)} parquet files")
+        
+        if not shard_files:
+            logger.warning("No parquet files found in shard directory")
+            
+        for f in sorted(shard_files):
+            try:
+                parts = f.stem.split("_")
+                if len(parts) < 3:
+                    logger.warning(f"Skipping file with invalid name format: {f}")
+                    continue
+                    
+                start = int(parts[1])
+                end = int(parts[2])
+                self.shard_map[(start, end)] = f.name
+                self.total_docs = max(self.total_docs, end + 1)
+                logger.debug(f"Mapped shard {f.name}: indices {start}-{end}")
+            except Exception as e:
+                logger.error(f"Error parsing shard filename {f}: {str(e)}")
+        
+        # Log shard statistics
+        logger.info(f"Shard map built with {len(self.shard_map)} shards")
+        logger.info(f"Total document count: {self.total_docs}")
+        
+        # Validate shard boundaries for gaps or overlaps
+        sorted_ranges = sorted(self.shard_map.keys())
+        for i in range(1, len(sorted_ranges)):
+            prev_end = sorted_ranges[i-1][1]
+            curr_start = sorted_ranges[i][0]
+            if curr_start != prev_end + 1:
+                logger.warning(f"Gap or overlap detected between shards: {prev_end} to {curr_start}")
 
     def get_metadata(self, global_indices):
         """Retrieve metadata with validation"""
         # Check for empty numpy array properly
         if isinstance(global_indices, np.ndarray) and global_indices.size == 0:
+            logger.warning("Empty indices array passed to get_metadata")
             return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
         
         # Convert numpy array to list for processing
         indices_list = global_indices.tolist() if isinstance(global_indices, np.ndarray) else global_indices
+        logger.info(f"Retrieving metadata for {len(indices_list)} indices")
         
         # Filter valid indices
         valid_indices = [idx for idx in indices_list if 0 <= idx < self.total_docs]
+        invalid_count = len(indices_list) - len(valid_indices)
+        if invalid_count > 0:
+            logger.warning(f"Filtered out {invalid_count} invalid indices")
+        
         if not valid_indices:
+            logger.warning("No valid indices remain after filtering")
             return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
     
         # Group indices by shard with boundary check
         shard_groups = {}
+        unassigned_indices = []
+        
         for idx in valid_indices:
             found = False
             for (start, end), shard in self.shard_map.items():
@@ -61,49 +123,137 @@ class MetadataManager:
                     found = True
                     break
             if not found:
-                st.warning(f"Index {idx} out of shard range (0-{self.total_docs-1})")
+                unassigned_indices.append(idx)
+                logger.warning(f"Index {idx} not found in any shard range")
+        
+        if unassigned_indices:
+            logger.warning(f"Could not assign {len(unassigned_indices)} indices to any shard")
     
         # Load and process shards
         results = []
         for shard, local_indices in shard_groups.items():
             try:
+                logger.info(f"Processing shard {shard} with {len(local_indices)} indices")
+                start_time = time.time()
+                
                 if shard not in self.loaded_shards:
-                    self.loaded_shards[shard] = pd.read_parquet(
-                        self.shard_dir / shard,
-                        columns=["title", "summary", "source"]
-                    )
+                    logger.info(f"Loading shard file: {shard}")
+                    shard_path = self.shard_dir / shard
+                    
+                    # Verify file exists
+                    if not shard_path.exists():
+                        logger.error(f"Shard file not found: {shard_path}")
+                        continue
+                        
+                    # Log file size
+                    file_size_mb = os.path.getsize(shard_path) / (1024 * 1024)
+                    logger.info(f"Shard file size: {file_size_mb:.2f} MB")
+                    
+                    # Attempt to read the parquet file
+                    try:
+                        self.loaded_shards[shard] = pd.read_parquet(
+                            shard_path,
+                            columns=["title", "summary", "source"]
+                        )
+                        logger.info(f"Successfully loaded shard {shard} with {len(self.loaded_shards[shard])} rows")
+                    except Exception as e:
+                        logger.error(f"Failed to read parquet file {shard}: {str(e)}")
+                        
+                        # Try to read file schema for debugging
+                        try:
+                            schema = pd.read_parquet(shard_path, engine='pyarrow').dtypes
+                            logger.info(f"Parquet schema: {schema}")
+                        except:
+                            pass
+                        continue
                 
                 if local_indices:
-                    results.append(self.loaded_shards[shard].iloc[local_indices])
+                    # Validate indices are within dataframe bounds
+                    df_len = len(self.loaded_shards[shard])
+                    valid_local_indices = [idx for idx in local_indices if 0 <= idx < df_len]
+                    
+                    if len(valid_local_indices) != len(local_indices):
+                        logger.warning(f"Filtered {len(local_indices) - len(valid_local_indices)} out-of-bounds indices")
+                    
+                    if valid_local_indices:
+                        logger.debug(f"Retrieving rows at indices: {valid_local_indices}")
+                        chunk = self.loaded_shards[shard].iloc[valid_local_indices]
+                        results.append(chunk)
+                        logger.info(f"Retrieved {len(chunk)} records from shard {shard}")
+                
+                logger.info(f"Shard processing completed in {time.time() - start_time:.2f} seconds")
+                    
             except Exception as e:
-                st.error(f"Error loading shard {shard}: {str(e)}")
+                logger.error(f"Error processing shard {shard}: {str(e)}", exc_info=True)
                 continue
     
-        return pd.concat(results).reset_index(drop=True) if results else pd.DataFrame()
+        # Combine results
+        if results:
+            combined = pd.concat(results).reset_index(drop=True)
+            logger.info(f"Combined metadata: {len(combined)} records from {len(results)} shards")
+            return combined
+        else:
+            logger.warning("No metadata records retrieved")
+            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
 
 class SemanticSearch:
     def __init__(self):
         self.shard_dir = Path("compressed_shards")
         self.model = None
         self.index_shards = []
-        self.metadata_mgr = MetadataManager()  # No Streamlit elements in constructor
+        self.metadata_mgr = MetadataManager()
         self.shard_sizes = []
         
+        # 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 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 = list(self.shard_dir.glob("*.index"))
+        self.logger.info(f"Found {len(index_files)} index files")
+        
         self.shard_sizes = []
-        for shard_path in sorted(self.shard_dir.glob("*.index")):
-            index = faiss.read_index(str(shard_path))
-            self.index_shards.append(index)
-            self.shard_sizes.append(index.ntotal)
+        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)}")
+        
+        total_vectors = sum(self.shard_sizes)
+        self.logger.info(f"Total loaded vectors: {total_vectors} across {len(self.index_shards)} shards")
 
     def _global_index(self, shard_idx, local_idx):
         """Convert local index to global index"""
@@ -111,67 +261,122 @@ class SemanticSearch:
 
     def search(self, query, top_k=5):
         """Search with validation"""
-        if not query or not self.index_shards:
+        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)
+            self.logger.debug(f"Query encoded to shape {query_embedding.shape}")
         except Exception as e:
-            st.error(f"Query encoding failed: {str(e)}")
+            self.logger.error(f"Query encoding failed: {str(e)}")
             return pd.DataFrame()
     
         all_distances = []
         all_global_indices = []
     
         # 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_indices = [idx for idx in indices[0] if 0 <= idx < index.ntotal]
+                
+                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(distances[0][:len(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:
-                st.error(f"Search failed in shard {shard_idx}: {str(e)}")
+                self.logger.error(f"Search failed in shard {shard_idx}: {str(e)}")
                 continue
     
-        # Ensure equal array lengths
-        min_length = min(len(all_distances), len(all_global_indices))
-        return self._process_results(
-            np.array(all_distances[:min_length]), 
-            np.array(all_global_indices[:min_length]), 
+        self.logger.info(f"Search found {len(all_global_indices)} results across all shards")
+        
+        # Process results
+        results = self._process_results(
+            np.array(all_distances), 
+            np.array(all_global_indices), 
             top_k
         )
+        
+        self.logger.info(f"Search completed in {time.time() - start_time:.2f} seconds with {len(results)} final results")
+        return results
 
     def _process_results(self, distances, global_indices, top_k):
         """Process raw search results into formatted DataFrame"""
+        process_start = time.time()
+        
         # 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", "similarity"])
         
         try:
-            # Convert numpy indices to Python list for metadata retrieval
-            indices_list = global_indices.tolist()
-            
             # Get metadata for matched indices
-            results = self.metadata_mgr.get_metadata(indices_list)
+            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")
             
+            # Empty results check
+            if len(results) == 0:
+                self.logger.warning("No metadata found for indices")
+                return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
+                
             # Ensure distances match results length
             if len(results) != len(distances):
-                distances = distances[:len(results)]
+                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)
             
         except Exception as e:
-            st.error(f"Result processing failed: {str(e)}")
-            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
+            self.logger.error(f"Result processing failed: {str(e)}", exc_info=True)
+            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])