Update app.py

app.py

@@ -83,7 +83,7 @@ def load_models():
         return False
 
 def load_embeddings():
-    """Load embeddings with robust error handling"""
+    """Load embeddings with robust error handling for numpy arrays"""
     try:
         print("Loading embeddings...")
         embeddings_path = 'embeddings.pkl'
@@ -91,33 +91,56 @@ def load_embeddings():
         if not os.path.exists(embeddings_path):
             print(f"Error: {embeddings_path} not found")
             return False
+
+        def persistent_load(pid):
+            return pid
             
         class CustomUnpickler(pickle.Unpickler):
+            def persistent_load(self, pid):
+                return pid
+                
             def find_class(self, module, name):
                 if module == "__main__":
                     module = "numpy"
                 return super().find_class(module, name)
-                
+        
         with open(embeddings_path, 'rb') as f:
             try:
-                data['embeddings'] = pickle.load(f)
+                # Try loading with numpy first
+                data['embeddings'] = np.load(f, allow_pickle=True).item()
             except Exception as e:
-                print(f"Standard unpickling failed, trying custom unpickler: {e}")
+                print(f"Numpy loading failed, trying pickle: {e}")
                 f.seek(0)
                 try:
-                    data['embeddings'] = CustomUnpickler(f).load()
+                    # Try standard pickle
+                    data['embeddings'] = pickle.load(f)
                 except Exception as e:
-                    print(f"Custom unpickler failed: {e}")
-                    data['embeddings'] = {}
-                    return False
+                    print(f"Standard pickle failed, trying custom unpickler: {e}")
+                    f.seek(0)
+                    try:
+                        # Try custom unpickler with persistent load handler
+                        unpickler = CustomUnpickler(f)
+                        data['embeddings'] = unpickler.load()
+                    except Exception as e:
+                        print(f"Custom unpickler failed: {e}")
+                        data['embeddings'] = {}
+                        return False
         
+        # Verify the loaded data
         if not isinstance(data['embeddings'], dict):
             print("Error: Embeddings data is not in expected format")
+            print(f"Actual type: {type(data['embeddings'])}")
             data['embeddings'] = {}
             return False
             
+        # Verify the structure of the embeddings
+        sample_key = next(iter(data['embeddings']))
+        sample_value = data['embeddings'][sample_key]
+        print(f"Sample embedding structure - Key: {sample_key}, Value type: {type(sample_value)}, Shape: {np.array(sample_value).shape}")
+        
         print(f"Successfully loaded {len(data['embeddings'])} embeddings")
         return True
+        
     except Exception as e:
         print(f"Error loading embeddings: {e}")
         data['embeddings'] = {}