Update search_utils.py

search_utils.py

@@ -9,10 +9,14 @@ import os
 import requests
 from functools import lru_cache
 from typing import List, Dict
+import pandas as pd
+from urllib.parse import quote
 
 # Configure logging
-logging.basicConfig(level=logging.WARNING,
-                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.WARNING,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger("OptimizedSearch")
 
 class OptimizedMetadataManager:
@@ -21,24 +25,27 @@ class OptimizedMetadataManager:
         self._init_url_resolver()
 
     def _init_metadata(self):
-        """Memory-mapped metadata loading"""
+        """Memory-mapped metadata loading.
+           Preloads all metadata (title and summary) into memory from parquet files.
+        """
         self.metadata_dir = Path("unzipped_cache/metadata_shards")
         self.metadata = {}
         
         # Preload all metadata into memory
         for parquet_file in self.metadata_dir.glob("*.parquet"):
             df = pd.read_parquet(parquet_file, columns=["title", "summary"])
+            # Using the dataframe index as key (assumes unique indices across files)
             self.metadata.update(df.to_dict(orient="index"))
             
         self.total_docs = len(self.metadata)
         logger.info(f"Loaded {self.total_docs} metadata entries into memory")
 
     def get_metadata_batch(self, indices: np.ndarray) -> List[Dict]:
-        """Batch retrieval of metadata"""
+        """Batch retrieval of metadata entries for a list of indices."""
         return [self.metadata.get(idx, {"title": "", "summary": ""}) for idx in indices]
 
     def _init_url_resolver(self):
-        """Initialize API session and cache"""
+        """Initialize API session and adapter for faster URL resolution."""
         self.session = requests.Session()
         adapter = requests.adapters.HTTPAdapter(
             pool_connections=10,
@@ -49,23 +56,26 @@ class OptimizedMetadataManager:
 
     @lru_cache(maxsize=10_000)
     def resolve_url(self, title: str) -> str:
-        """Optimized URL resolution with fail-fast"""
+        """Optimized URL resolution with caching and a fail-fast approach."""
         try:
             # Try arXiv first
             arxiv_url = self._get_arxiv_url(title)
-            if arxiv_url: return arxiv_url
+            if arxiv_url:
+                return arxiv_url
             
             # Fallback to Semantic Scholar
             semantic_url = self._get_semantic_url(title)
-            if semantic_url: return semantic_url
+            if semantic_url:
+                return semantic_url
             
         except Exception as e:
             logger.warning(f"URL resolution failed: {str(e)}")
             
+        # Default fallback to Google Scholar search
         return f"https://scholar.google.com/scholar?q={quote(title)}"
 
     def _get_arxiv_url(self, title: str) -> str:
-        """Fast arXiv lookup with timeout"""
+        """Fast arXiv lookup with a short timeout."""
         with self.session.get(
             "http://export.arxiv.org/api/query",
             params={"search_query": f'ti:"{title}"', "max_results": 1},
@@ -76,14 +86,15 @@ class OptimizedMetadataManager:
         return ""
 
     def _parse_arxiv_response(self, xml: str) -> str:
-        """Fast XML parsing using string operations"""
-        if "<entry>" not in xml: return ""
+        """Fast XML parsing using simple string operations."""
+        if "<entry>" not in xml:
+            return ""
         start = xml.find("<id>") + 4
         end = xml.find("</id>", start)
         return xml[start:end].replace("http:", "https:") if start > 3 else ""
 
     def _get_semantic_url(self, title: str) -> str:
-        """Batch-friendly Semantic Scholar lookup"""
+        """Semantic Scholar lookup with a short timeout."""
         with self.session.get(
             "https://api.semanticscholar.org/graph/v1/paper/search",
             params={"query": title[:200], "limit": 1},
@@ -97,53 +108,62 @@ class OptimizedMetadataManager:
 
 class OptimizedSemanticSearch:
     def __init__(self):
+        # Load the sentence transformer model
         self.model = SentenceTransformer('all-MiniLM-L6-v2')
         self._load_faiss_indexes()
         self.metadata_mgr = OptimizedMetadataManager()
         
     def _load_faiss_indexes(self):
-        """Load indexes with memory mapping"""
+        """Load the FAISS index with memory mapping for read-only access."""
+        # Here we assume the FAISS index has been combined into one file.
         self.index = faiss.read_index("combined_index.faiss", faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)
         logger.info(f"Loaded FAISS index with {self.index.ntotal} vectors")
 
     def search(self, query: str, top_k: int = 5) -> List[Dict]:
-        """Optimized search pipeline"""
+        """Optimized search pipeline:
+           - Encodes the query.
+           - Performs FAISS search (fetching extra results for deduplication).
+           - Retrieves metadata and processes results.
+        """
         # Batch encode query
         query_embedding = self.model.encode([query], convert_to_numpy=True)
         
-        # FAISS search
-        distances, indices = self.index.search(query_embedding, top_k*2)  # Search extra for dedup
+        # FAISS search: we search for more than top_k to allow for deduplication.
+        distances, indices = self.index.search(query_embedding, top_k * 2)
         
         # Batch metadata retrieval
         results = self.metadata_mgr.get_metadata_batch(indices[0])
         
-        # Process results
+        # Process and return the final results
         return self._process_results(results, distances[0], top_k)
 
     def _process_results(self, results: List[Dict], distances: np.ndarray, top_k: int) -> List[Dict]:
-        """Parallel result processing"""
+        """Parallel processing of search results:
+           - Resolve source URLs in parallel.
+           - Add similarity scores.
+           - Deduplicate and sort the results.
+        """
         with concurrent.futures.ThreadPoolExecutor() as executor:
-            # Parallel URL resolution
+            # Parallel URL resolution for each result
             futures = {
                 executor.submit(
                     self.metadata_mgr.resolve_url, 
                     res["title"]
                 ): idx for idx, res in enumerate(results)
             }
-            
-            # Update results as URLs resolve
+            # Update each result as URLs resolve
             for future in concurrent.futures.as_completed(futures):
                 idx = futures[future]
                 try:
                     results[idx]["source"] = future.result()
                 except Exception as e:
                     results[idx]["source"] = ""
-
-        # Add similarity scores
+        
+        # Add similarity scores based on distances
         for idx, dist in enumerate(distances[:len(results)]):
             results[idx]["similarity"] = 1 - (dist / 2)
 
-        # Deduplicate and sort
+        # Deduplicate by title and sort by similarity score (descending)
         seen = set()
         final_results = []
         for res in sorted(results, key=lambda x: x["similarity"], reverse=True):