Update search_utils.py

search_utils.py

@@ -4,6 +4,62 @@ import faiss
 from pathlib import Path
 from sentence_transformers import SentenceTransformer, util
 import streamlit as st
+import zipfile
+import pandas as pd
+from pathlib import Path
+import streamlit as st
+
+class MetadataManager:
+    def __init__(self):
+        self.shard_dir = Path("metadata_shards")
+        self.shard_map = {}
+        self.loaded_shards = {}
+        self._ensure_unzipped()
+        self._build_shard_map()
+
+    def _ensure_unzipped(self):
+        """Extract metadata shards from zip if needed"""
+        if not self.shard_dir.exists():
+            zip_path = Path("metadata_shards.zip")
+            if zip_path.exists():
+                with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+                    zip_ref.extractall(self.shard_dir)
+                st.toast("✅ Successfully extracted metadata shards!", icon="📦")
+            else:
+                raise FileNotFoundError("No metadata shards found!")
+
+    def _build_shard_map(self):
+        """Map index ranges to shard files"""
+        for f in self.shard_dir.glob("*.parquet"):
+            parts = f.stem.split("_")
+            self.shard_map[(int(parts[1]), int(parts[2]))] = f.name
+
+    def get_metadata(self, indices):
+        """Retrieve metadata for specific indices"""
+        results = []
+        shard_groups = {}
+        
+        # Group indices by shard
+        for idx in indices:
+            for (start, end), shard in self.shard_map.items():
+                if start <= idx <= end:
+                    if shard not in shard_groups:
+                        shard_groups[shard] = []
+                    shard_groups[shard].append(idx - start)
+                    break
+        
+        # Load required shards
+        for shard, local_indices in shard_groups.items():
+            if shard not in self.loaded_shards:
+                self.loaded_shards[shard] = pd.read_parquet(
+                    self.shard_dir / shard,
+                    columns=["title", "summary", "source"]
+                )
+            
+            results.append(self.loaded_shards[shard].iloc[local_indices])
+        
+        return pd.concat(results).reset_index(drop=True)
+
 
 class SemanticSearch:
     def __init__(self, shard_dir="compressed_shards"):
@@ -11,6 +67,7 @@ class SemanticSearch:
         self.shard_dir.mkdir(exist_ok=True, parents=True)
         self.model = None
         self.index_shards = []
+        self.metadata_mgr = MetadataManager()
         
     @st.cache_resource
     def load_model(_self):
@@ -61,4 +118,55 @@ class SemanticSearch:
         """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)
+        return filtered.reset_index(drop=True)
+
+
+
+
+class MetadataManager:
+    def __init__(self, repo_id, shard_dir="metadata_shards"):
+        self.repo_id = repo_id
+        self.shard_dir = Path(shard_dir)
+        self.shard_map = {}
+        self.loaded_shards = {}
+        self._build_shard_map()
+
+    def _build_shard_map(self):
+        """Map index ranges to shard files"""
+        for f in self.shard_dir.glob("*.parquet"):
+            parts = f.stem.split("_")
+            self.shard_map[(int(parts[1]), int(parts[2]))] = f.name
+
+    def _download_shard(self, shard_name):
+        """Download missing shards on demand"""
+        if not (self.shard_dir/shard_name).exists():
+            hf_hub_download(
+                repo_id=self.repo_id,
+                filename=f"metadata_shards/{shard_name}",
+                local_dir=self.shard_dir,
+                cache_dir="metadata_cache"
+            )
+
+    def get_metadata(self, indices):
+        """Retrieve metadata for specific indices"""
+        results = []
+        
+        # Group indices by shard
+        shard_groups = {}
+        for idx in indices:
+            for (start, end), shard in self.shard_map.items():
+                if start <= idx <= end:
+                    if shard not in shard_groups:
+                        shard_groups[shard] = []
+                    shard_groups[shard].append(idx - start)
+                    break
+        
+        # Process each required shard
+        for shard, local_indices in shard_groups.items():
+            if shard not in self.loaded_shards:
+                self._download_shard(shard)
+                self.loaded_shards[shard] = pd.read_parquet(self.shard_dir/shard)
+            
+            results.append(self.loaded_shards[shard].iloc[local_indices])
+        
+        return pd.concat(results).reset_index(drop=True)