Update app.py

app.py

@@ -87,9 +87,14 @@ import numpy as np
 import os
 from typing import Dict, Optional
 
+class EmbeddingsUnpickler(pickle.Unpickler):
+    def persistent_load(self, pid):
+        # Handle persistent IDs by returning them as-is
+        return pid
+
 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.
+    Load embeddings from a pickle file with support for persistent IDs.
     
     Args:
         embeddings_path (str): Path to the pickle file containing embeddings
@@ -103,7 +108,9 @@ def load_embeddings(embeddings_path: str = 'embeddings.pkl') -> Optional[Dict[st
         
     try:
         with open(embeddings_path, 'rb') as f:
-            embeddings = pickle.load(f)
+            # Use custom unpickler with persistent_load support
+            unpickler = EmbeddingsUnpickler(f)
+            embeddings = unpickler.load()
             
         # Validate the loaded data
         if not isinstance(embeddings, dict):
@@ -111,18 +118,30 @@ def load_embeddings(embeddings_path: str = 'embeddings.pkl') -> Optional[Dict[st
             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])
+        processed_embeddings = {}
+        for key, value in embeddings.items():
+            # Handle both direct arrays and persistent IDs
+            if isinstance(value, (list, np.ndarray)):
+                processed_embeddings[key] = np.array(value)
+            else:
+                # If it's a persistent ID, convert it to a numpy array
+                try:
+                    processed_embeddings[key] = np.array(value)
+                except Exception as e:
+                    print(f"Warning: Could not convert embedding for {key}: {e}")
+                    continue
         
         # 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
+        if processed_embeddings:
+            sample_key = next(iter(processed_embeddings))
+            print(f"Data type: {type(processed_embeddings)}")
+            print(f"First few keys and values:")
+            print(f"Key: {sample_key}, Value: {processed_embeddings[sample_key][:20]}")
+            print(f"Successfully loaded {len(processed_embeddings)} embeddings")
+            return processed_embeddings
+        else:
+            print("Error: No valid embeddings were processed")
+            return None
         
     except Exception as e:
         print(f"Error loading embeddings: {str(e)}")