Update app.py

app.py

@@ -25,6 +25,7 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file 
 from typing import List, Dict, Optional
 from safetensors.numpy import load_file
+from safetensors.torch import safe_open
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -112,7 +113,6 @@ def load_models():
         return False
 
 def load_embeddings() -> Optional[Dict[str, np.ndarray]]:
-    """Load embeddings from Safetensors file"""
     try:
         # Locate or download embeddings file
         embeddings_path = 'embeddings.safetensors'
@@ -124,25 +124,38 @@ def load_embeddings() -> Optional[Dict[str, np.ndarray]]:
                 repo_type="space"
             )
         
-        # Load Safetensors file
-        embeddings = load_file(embeddings_path)
-        if not isinstance(embeddings, dict):
-            raise ValueError("Expected a dictionary in the Safetensors file.")
-        
-        # Validate and convert tensors to numpy arrays
-        result = {}
-        for key, tensor in embeddings.items():
-            if not hasattr(tensor, 'numpy'):
-                raise TypeError(f"Value for key {key} is not a tensor or cannot be converted to numpy.")
-            result[key] = tensor.numpy()
+        # Initialize a dictionary to store embeddings
+        embeddings = {}
+
+        # Open the safetensors file
+        with safe_open(embeddings_path, framework="pt") as f:
+            keys = f.keys()
+            print(f"Available keys in the .safetensors file: {list(keys)}")  # Debugging info
+
+            # Iterate over the keys and load tensors
+            for key in keys:
+                try:
+                    tensor = f.get_tensor(key)
+                    if not isinstance(tensor, torch.Tensor):
+                        raise TypeError(f"Value for key {key} is not a valid PyTorch tensor.")
+                    
+                    # Convert tensor to NumPy array
+                    embeddings[key] = tensor.numpy()
+                except Exception as key_error:
+                    print(f"Failed to process key {key}: {key_error}")
+
+        if embeddings:
+            print("Embeddings successfully loaded.")
+        else:
+            print("No embeddings could be loaded. Please check the file format and content.")
         
-        print("Embeddings successfully loaded.")
-        return result
+        return embeddings
 
     except Exception as e:
         print(f"Error loading embeddings: {e}")
         return None
 
+
 def load_documents_data(folder_path='downloaded_articles/downloaded_articles'):
     """Load document data from HTML articles in a specified folder."""
     try: