datamodule: store num_of_labels to hparms

Author: aditya0by0

chebai/preprocessing/datasets/base.py

@@ -401,6 +401,21 @@ def setup(self, **kwargs):
         if not ("keep_reader" in kwargs and kwargs["keep_reader"]):
             self.reader.on_finish()
 
+        self._add_num_of_labels_to_hparams()
+
+    def _add_num_of_labels_to_hparams(self):
+        num_of_labels = len(
+            torch.load(
+                os.path.join(
+                    self.processed_dir, self.processed_file_names_dict["data"]
+                ),
+                weights_only=False,
+            )[0]["labels"]
+        )
+
+        print(f"Number of labels for loaded data: {num_of_labels}")
+        self.hparams.num_of_labels = num_of_labels
+
     def setup_processed(self):
         """
         Setup the processed data.
@@ -541,6 +556,8 @@ def setup(self, **kwargs):
         for s in self.subsets:
             s.setup(**kwargs)
 
+        self._add_num_of_labels_to_hparams()
+
     def dataloader(self, kind: str, **kwargs) -> DataLoader:
         """
         Creates a DataLoader for a specific subset.

chebai/preprocessing/datasets/tox21.py

@@ -129,6 +129,8 @@ def setup(self, **kwargs) -> None:
         ):
             self.setup_processed()
 
+        self._add_num_of_labels_to_hparams()
+
     def _load_data_from_file(self, input_file_path: str) -> List[Dict]:
         """Loads data from a CSV file.
 
@@ -311,6 +313,8 @@ def setup(self, **kwargs) -> None:
         ):
             self.setup_processed()
 
+        self._add_num_of_labels_to_hparams()
+
     def _load_dict(self, input_file_path: str) -> Generator[Dict, None, None]:
         """Loads data from a CSV file as a generator.