Update app.py

app.py

@@ -22,38 +22,74 @@ import time
 # Initialize FastAPI app first
 app = FastAPI()
 
-class CustomUnpickler(pickle.Unpickler):
+class ArticleEmbeddingUnpickler(pickle.Unpickler):
+    """Custom unpickler specifically designed for article embeddings"""
+    def find_class(self, module, name):
+        # Handle numpy arrays specially
+        if module == 'numpy':
+            return getattr(np, name)
+        # Handle the SentenceTransformer case
+        if module == 'sentence_transformers.SentenceTransformer':
+            from sentence_transformers import SentenceTransformer
+            return SentenceTransformer
+        return super().find_class(module, name)
+    
     def persistent_load(self, pid):
+        """Handle persistent IDs during unpickling"""
         try:
+            # Convert to string if bytes
             if isinstance(pid, bytes):
                 pid = pid.decode('utf-8', errors='ignore')
-            pid = str(pid).encode('ascii', errors='ignore').decode('ascii')
-            
-            if pid == "sentence_transformer_model":
-                return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-            return pid
+            return str(pid)
         except Exception as e:
-            raise pickle.UnpicklingError(f"Error handling persistent ID: {e}")
+            print(f"Error in persistent_load: {str(e)}")
+            return str(pid)
 
-def safe_load_embeddings():
+def safe_load_embeddings(file_path='embeddings.pkl'):
+    """Load embeddings with enhanced error handling for article embeddings"""
     try:
-        with open('embeddings.pkl', 'rb') as file:
-            unpickler = CustomUnpickler(file)
+        if not os.path.exists(file_path):
+            print(f"Embeddings file not found at {file_path}")
+            return None
+
+        with open(file_path, 'rb') as file:
+            unpickler = ArticleEmbeddingUnpickler(file)
             embeddings_data = unpickler.load()
-            
+
+            # Validate the dictionary structure
             if not isinstance(embeddings_data, dict):
-                raise ValueError("Loaded data is not a dictionary")
-                
-            first_key = next(iter(embeddings_data))
-            if not isinstance(embeddings_data[first_key], (np.ndarray, list)):
-                raise ValueError("Embeddings are not in the expected format")
-                
-            return embeddings_data
-            
-    except (FileNotFoundError, pickle.UnpicklingError, ValueError) as e:
+                print(f"Invalid data structure: expected dict, got {type(embeddings_data)}")
+                return None
+
+            # Validate each embedding
+            valid_embeddings = {}
+            for key, value in embeddings_data.items():
+                # Ensure key is string and value is numpy array
+                try:
+                    key_str = str(key)
+                    if isinstance(value, list):
+                        value = np.array(value)
+                    if isinstance(value, np.ndarray):
+                        valid_embeddings[key_str] = value
+                    else:
+                        print(f"Skipping invalid embedding for key {key}: {type(value)}")
+                except Exception as e:
+                    print(f"Error processing embedding {key}: {str(e)}")
+                    continue
+
+            if not valid_embeddings:
+                print("No valid embeddings found in file")
+                return None
+
+            print(f"Successfully loaded {len(valid_embeddings)} embeddings")
+            return valid_embeddings
+
+    except Exception as e:
         print(f"Error loading embeddings: {str(e)}")
         return None
 
+
+
 # Models and data structures
 class GlobalModels:
     embedding_model = None
@@ -91,10 +127,18 @@ class DocumentResponse(BaseModel):
     text: str
     score: float
 
+# Modified startup event handler
 @app.on_event("startup")
 async def load_models():
-    """Initialize all models and data on startup"""
+    """Initialize models with enhanced embeddings loading"""
     try:
+        # Load embeddings first
+        embeddings_data = safe_load_embeddings()
+        if embeddings_data is None:
+            raise HTTPException(
+                status_code=500,
+                detail="Failed to load embeddings data. Check logs for details."
+            )
         # Load embedding models first
         global_models.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')