Skip to content

Commit a21f617

Browse files
committed
after_setup method for processing internal prop after setup
1 parent d0ed76a commit a21f617

File tree

1 file changed

+14
-2
lines changed
  • chebai/preprocessing/datasets

1 file changed

+14
-2
lines changed

chebai/preprocessing/datasets/base.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -440,13 +440,25 @@ def setup(self, *args, **kwargs) -> None:
440440
):
441441
self.setup_processed()
442442

443-
if not ("keep_reader" in kwargs and kwargs["keep_reader"]):
444-
self.reader.on_finish()
443+
self._after_setup(**kwargs)
445444

445+
def _after_setup(self, **kwargs):
446+
"""
447+
Finalize the setup process after ensuring the processed data is available.
448+
449+
This method performs post-setup tasks like finalizing the reader and setting internal properties.
450+
"""
451+
self.reader.on_finish()
446452
self._set_processed_data_props()
447453

448454
def _set_processed_data_props(self):
455+
"""
456+
Load processed data and extract metadata.
449457
458+
Sets:
459+
- self._num_of_labels: Number of target labels in the dataset.
460+
- self._feature_vector_size: Maximum feature vector length across all data points.
461+
"""
450462
data_pt = torch.load(
451463
os.path.join(self.processed_dir, self.processed_file_names_dict["data"]),
452464
weights_only=False,

0 commit comments

Comments
 (0)