From a21f617f95a9afee82612c4cd2034b9ecfbc68c0 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 31 May 2025 20:09:08 +0200 Subject: [PATCH] after_setup method for processing internal prop after setup --- chebai/preprocessing/datasets/base.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 40623dd2..6698ab8d 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -440,13 +440,25 @@ def setup(self, *args, **kwargs) -> None: ): self.setup_processed() - if not ("keep_reader" in kwargs and kwargs["keep_reader"]): - self.reader.on_finish() + self._after_setup(**kwargs) + def _after_setup(self, **kwargs): + """ + Finalize the setup process after ensuring the processed data is available. + + This method performs post-setup tasks like finalizing the reader and setting internal properties. + """ + self.reader.on_finish() self._set_processed_data_props() def _set_processed_data_props(self): + """ + Load processed data and extract metadata. + Sets: + - self._num_of_labels: Number of target labels in the dataset. + - self._feature_vector_size: Maximum feature vector length across all data points. + """ data_pt = torch.load( os.path.join(self.processed_dir, self.processed_file_names_dict["data"]), weights_only=False,