add prepare_data method to base class

- after doing certain common processing the  `prepare_data` method calls the _perform_data_preparation method which contain core logic for data preparation

Author: aditya0by0

chebai/preprocessing/datasets/base.py

@@ -138,9 +138,10 @@ def __init_subclass__(cls, *args, **kwargs):
         super().__init_subclass__(*args, **kwargs)
         original_init = cls.__init__
 
+        # Creates updated definition for init method
         def new_init(self, *args, **kwargs):
             original_init(self, *args, **kwargs)  # Call the original __init__
-            if type(self) == cls:  # Only run __post_init__ if it's the final class
+            if type(self) == cls:  # Only run method if it's the final class
                 self._call_data_processing_methods(*args, **kwargs)
 
         cls.__init__ = new_init
@@ -448,6 +449,10 @@ def prepare_data(self, *args, **kwargs) -> None:
             return
 
         self._prepare_data_flag += 1
+        self._perform_data_preparation(*args, **kwargs)
+
+    def _perform_data_preparation(self, *args, **kwargs) -> None:
+        raise NotImplementedError
 
     def setup(self, *args, **kwargs) -> None:
         """
@@ -598,11 +603,10 @@ def __init__(
         os.makedirs(self.processed_dir, exist_ok=True)
         super(pl.LightningDataModule, self).__init__(**kwargs)
 
-    def prepare_data(self):
+    def _perform_data_preparation(self):
         """
         Placeholder for data preparation logic.
         """
-        super().prepare_data()
         for s in self.subsets:
             s.prepare_data()
 
@@ -792,7 +796,7 @@ def _validate_splits_file_path(splits_file_path: Optional[str]) -> Optional[str]
         return splits_file_path
 
     # ------------------------------ Phase: Prepare data -----------------------------------
-    def prepare_data(self, *args: Any, **kwargs: Any) -> None:
+    def _perform_data_preparation(self, *args: Any, **kwargs: Any) -> None:
         """
         Prepares the data for the dataset.
 
@@ -811,7 +815,6 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
         Returns:
             None
         """
-        super().prepare_data()
         print("Checking for processed data in", self.processed_dir_main)
 
         processed_name = self.processed_main_file_names_dict["data"]

chebai/preprocessing/datasets/chebi.py

@@ -59,8 +59,7 @@ def download(self):
     def raw_file_names(self):
         return ["test.pkl", "train.pkl", "validation.pkl"]
 
-    def prepare_data(self, *args, **kwargs):
-        super().prepare_data()
+    def _perform_data_preparation(self, *args, **kwargs):
         print("Check for raw data in", self.raw_dir)
         if any(
             not os.path.isfile(os.path.join(self.raw_dir, f))
@@ -156,7 +155,7 @@ def __init__(
             )
 
     # ------------------------------ Phase: Prepare data -----------------------------------
-    def prepare_data(self, *args: Any, **kwargs: Any) -> None:
+    def _perform_data_preparation(self, *args: Any, **kwargs: Any) -> None:
         """
         Prepares the data for the Chebi dataset.
 
@@ -177,8 +176,6 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
         Returns:
             None
         """
-        super().prepare_data(args, kwargs)
-
         if self.chebi_version_train is not None:
             if not os.path.isfile(
                 os.path.join(

chebai/preprocessing/datasets/deepGO/go_uniprot.py

@@ -770,7 +770,7 @@ def _name(self) -> str:
         return f"{threshold_part}{self.max_sequence_length}"
 
     # ------------------------------ Phase: Prepare data -----------------------------------
-    def prepare_data(self, *args: Any, **kwargs: Any) -> None:
+    def _perform_data_preparation(self, *args: Any, **kwargs: Any) -> None:
         """
         Checks for the existence of migrated DeepGO data in the specified directory.
         Raises an error if the required data file is not found, prompting
@@ -783,11 +783,6 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
         Raises:
             FileNotFoundError: If the processed data file does not exist.
         """
-        if self._prepare_data_flag != 1:
-            return
-
-        self._prepare_data_flag += 1
-
         print("Checking for processed data in", self.processed_dir_main)
 
         processed_name = self.processed_main_file_names_dict["data"]

chebai/preprocessing/datasets/deepGO/protein_pretraining.py

@@ -55,7 +55,7 @@ def __init__(self, **kwargs):
             )
 
     # ------------------------------ Phase: Prepare data -----------------------------------
-    def prepare_data(self, *args: Any, **kwargs: Any) -> None:
+    def _perform_data_preparation(self, *args: Any, **kwargs: Any) -> None:
         """
         Prepares the data by downloading and parsing Swiss-Prot data if not already available. Saves the processed data
         for further use.
@@ -64,11 +64,6 @@ def prepare_data(self, *args: Any, **kwargs: Any) -> None:
             *args: Additional positional arguments.
             **kwargs: Additional keyword arguments.
         """
-        if self._prepare_data_flag != 1:
-            return
-
-        self._prepare_data_flag += 1
-
         processed_name = self.processed_main_file_names_dict["data"]
         if not os.path.isfile(os.path.join(self.processed_dir_main, processed_name)):
             print("Missing processed data file (`data.pkl` file)")

chebai/preprocessing/datasets/pubchem.py

@@ -179,11 +179,10 @@ def processed_file_names(self) -> List[str]:
         """
         return ["test.pt", "train.pt", "validation.pt"]
 
-    def prepare_data(self, *args, **kwargs):
+    def _perform_data_preparation(self, *args, **kwargs):
         """
         Checks for raw data and downloads if necessary.
         """
-        super().prepare_data()
         print("Check for raw data in", self.raw_dir)
         if any(
             not os.path.isfile(os.path.join(self.raw_dir, f))