Update app.py

app.py

@@ -91,25 +91,29 @@ def load_models():
 def load_embeddings() -> Optional[Dict[str, np.ndarray]]:
     """Load embeddings from local file or HuggingFace Hub"""
     try:
-        # First try local file
+        import pickle
+        import numpy as np
+        import os
+        from typing import Dict, Optional
+
         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
+        if not os.path.exists(embeddings_path):
             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
+            embeddings_path = hf_hub_download(
+                repo_id=os.environ.get('HF_SPACE_ID', ''),
                 filename="embeddings.pkl",
                 repo_type="space"
             )
-            with open(file_path, 'rb') as f:
-                unpickler = pickle.Unpickler(f)
-                unpickler.encoding = 'ascii'
-                embeddings = unpickler.load()
+
+        class ASCIIUnpickler(pickle.Unpickler):
+            def find_class(self, module, name):
+                if module == "__main__":
+                    module = "numpy"
+                return super().find_class(module, name)
+
+        with open(embeddings_path, 'rb') as f:
+            unpickler = ASCIIUnpickler(f)
+            embeddings = unpickler.load()
 
         if not isinstance(embeddings, dict):
             return None