Update app.py

app.py

@@ -94,32 +94,38 @@ def load_models():
         return False
 
 
-def load_embeddings(repo_id: str) -> Optional[Dict[str, np.ndarray]]:
-    """Load embeddings using HuggingFace Hub"""
+def load_embeddings() -> Optional[Dict[str, np.ndarray]]:
+    """Load embeddings from local file or HuggingFace Hub"""
     try:
-        # Download file from HF Hub
-        file_path = hf_hub_download(
-            repo_id=repo_id,
-            filename="embeddings.pkl",
-            repo_type="space"
-        )
-        
-        # Load with custom unpickler
-        with open(file_path, 'rb') as f:
-            unpickler = pickle.Unpickler(f)
-            unpickler.encoding = 'ascii'
-            embeddings = unpickler.load()
+        # First try local file
+        embeddings_path = 'embeddings.pkl'
+        if os.path.exists(embeddings_path):
+            with open(embeddings_path, 'rb') as f:
+                unpickler = pickle.Unpickler(f)
+                unpickler.encoding = 'ascii'
+                embeddings = unpickler.load()
+        else:
+            # If local file doesn't exist, try downloading from HF Hub
+            from huggingface_hub import hf_hub_download
+            file_path = hf_hub_download(
+                repo_id=os.environ.get('HF_SPACE_ID', ''),  # Gets Space ID from environment
+                filename="embeddings.pkl",
+                repo_type="space"
+            )
+            with open(file_path, 'rb') as f:
+                unpickler = pickle.Unpickler(f)
+                unpickler.encoding = 'ascii'
+                embeddings = unpickler.load()
+
+        if not isinstance(embeddings, dict):
+            return None
             
-            if not isinstance(embeddings, dict):
-                return None
-                
-            # Convert to numpy arrays
-            return {k: np.array(v, dtype=np.float32) for k, v in embeddings.items()}
+        return {k: np.array(v, dtype=np.float32) for k, v in embeddings.items()}
             
     except Exception as e:
         print(f"Error loading embeddings: {e}")
         return None
-
+        
 def load_documents_data():
     """Load document data with error handling"""
     try: