separate load_processed_data into file-based and kind-based loading

Author: sfluegel05

chebai/preprocessing/datasets/base.py

@@ -1103,8 +1103,8 @@ def _retrieve_splits_from_csv(self) -> None:
         splits_df = pd.read_csv(self.splits_file_path)
 
         filename = self.processed_file_names_dict["data"]
-        data = torch.load(
-            os.path.join(self.processed_dir, filename), weights_only=False
+        data = self.load_processed_data_from_file(
+            os.path.join(self.processed_dir, filename)
         )
         df_data = pd.DataFrame(data)
 
@@ -1161,9 +1161,10 @@ def load_processed_data(
                 return data_df.to_dict(orient="records")
 
         # If filename is provided
-        return torch.load(
-            os.path.join(self.processed_dir, filename), weights_only=False
-        )
+        return self.load_processed_data_from_file(filename)
+
+    def load_processed_data_from_file(self, filename):
+        return torch.load(os.path.join(filename), weights_only=False)
 
     # ------------------------------ Phase: Raw Properties -----------------------------------
     @property

chebai/preprocessing/datasets/chebi.py

@@ -401,8 +401,8 @@ def _get_data_splits(self) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
         """
         try:
             filename = self.processed_file_names_dict["data"]
-            data_chebi_version = torch.load(
-                os.path.join(self.processed_dir, filename), weights_only=False
+            data_chebi_version = self.load_processed_data_from_file(
+                os.path.join(self.processed_dir, filename)
             )
         except FileNotFoundError:
             raise FileNotFoundError(