Update app.py

app.py

@@ -82,69 +82,51 @@ def load_models():
         print(f"Error loading models: {e}")
         return False
 
-def load_embeddings():
-    """Load embeddings with robust error handling for numpy arrays"""
-    try:
-        print("Loading embeddings...")
-        embeddings_path = 'embeddings.pkl'
-        
-        if not os.path.exists(embeddings_path):
-            print(f"Error: {embeddings_path} not found")
-            return False
+import pickle
+import numpy as np
+import os
+from typing import Dict, Optional
 
-        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)
+def load_embeddings(embeddings_path: str = 'embeddings.pkl') -> Optional[Dict[str, np.ndarray]]:
+    """
+    Load embeddings from a pickle file containing a dictionary of numpy arrays.
+    
+    Args:
+        embeddings_path (str): Path to the pickle file containing embeddings
         
-        with open(embeddings_path, 'rb') as f:
-            try:
-                # Try loading with numpy first
-                data['embeddings'] = np.load(f, allow_pickle=True).item()
-            except Exception as e:
-                print(f"Numpy loading failed, trying pickle: {e}")
-                f.seek(0)
-                try:
-                    # Try standard pickle
-                    data['embeddings'] = pickle.load(f)
-                except Exception as e:
-                    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
+    Returns:
+        Optional[Dict[str, np.ndarray]]: Dictionary of embeddings or None if loading fails
+    """
+    if not os.path.exists(embeddings_path):
+        print(f"Error: {embeddings_path} not found")
+        return None
         
-        # 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
+    try:
+        with open(embeddings_path, 'rb') as f:
+            embeddings = pickle.load(f)
             
-        # 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}")
+        # Validate the loaded data
+        if not isinstance(embeddings, dict):
+            print(f"Error: Expected dict, got {type(embeddings)}")
+            return None
+            
+        # Convert values to numpy arrays if they aren't already
+        for key in embeddings:
+            if not isinstance(embeddings[key], np.ndarray):
+                embeddings[key] = np.array(embeddings[key])
         
-        print(f"Successfully loaded {len(data['embeddings'])} embeddings")
-        return True
+        # Print sample for verification
+        sample_key = next(iter(embeddings))
+        print(f"Data type: {type(embeddings)}")
+        print(f"First few keys and values:")
+        print(f"Key: {sample_key}, Value: {embeddings[sample_key][:20]}")  # Show first 20 values
+        print(f"Successfully loaded {len(embeddings)} embeddings")
+        
+        return embeddings
         
     except Exception as e:
-        print(f"Error loading embeddings: {e}")
-        data['embeddings'] = {}
-        return False
+        print(f"Error loading embeddings: {str(e)}")
+        return None
 
 def load_documents_data():
     """Load document data with error handling"""