Fixed code on loading cell embeddings for package

protac_degradation_predictor/data_utils.py

@@ -19,24 +19,25 @@ memory = Memory(cachedir, verbose=0)
 
 
 @memory.cache
-def load_cell2embedding(
-        embeddings_path: Optional[str] = None,
+def load_protein2embedding(
+    embeddings_path: Optional[str] = None,
 ) -> Dict[str, np.ndarray]:
-    """ Load the cell line embeddings from a file.
-    
+    """ Load the protein embeddings from a file.
+
     Args:
         embeddings_path (str): The path to the embeddings file.
-        
+
     Returns:
-        Dict[str, np.ndarray]: A dictionary of cell line embeddings.
+        Dict[str, np.ndarray]: A dictionary of protein embeddings.
     """
     if embeddings_path is None:
-        with pkg_resources.resource_stream(__name__, 'data/cell2embedding.pkl') as f:
-            cell2embedding = pickle.load(f)
-    else:
-        with open(embeddings_path, 'rb') as f:
-            cell2embedding = pickle.load(f)
-    return cell2embedding
+        embeddings_path = pkg_resources.resource_stream(__name__, 'data/uniprot2embedding.h5')
+    protein2embedding = {}
+    with h5py.File(embeddings_path, "r") as file:
+        for sequence_id in file.keys():
+            embedding = file[sequence_id][:]
+            protein2embedding[sequence_id] = np.array(embedding)
+    return protein2embedding
 
 
 @memory.cache
@@ -52,9 +53,11 @@ def load_cell2embedding(
         Dict[str, np.ndarray]: A dictionary of cell line embeddings.
     """
     if embeddings_path is None:
-        embeddings_path = pkg_resources.resource_stream(__name__, 'data/cell2embedding.pkl')
-    with open(embeddings_path, 'rb') as f:
-        cell2embedding = pickle.load(f)
+        with pkg_resources.resource_stream(__name__, 'data/cell2embedding.pkl') as f:
+            cell2embedding = pickle.load(f)
+    else:
+        with open(embeddings_path, 'rb') as f:
+            cell2embedding = pickle.load(f)
     return cell2embedding