add docstring

Author: aditya0by0

chebai/preprocessing/datasets/base.py

@@ -96,9 +96,9 @@ def __init__(
         self.prediction_kind = prediction_kind
         self.data_limit = data_limit
         self.label_filter = label_filter
-        assert (balance_after_filter is not None) or (
-            self.label_filter is None
-        ), "Filter balancing requires a filter"
+        assert (balance_after_filter is not None) or (self.label_filter is None), (
+            "Filter balancing requires a filter"
+        )
         self.balance_after_filter = balance_after_filter
         self.num_workers = num_workers
         self.persistent_workers: bool = bool(persistent_workers)
@@ -108,13 +108,13 @@ def __init__(
         self.use_inner_cross_validation = (
             inner_k_folds > 1
         )  # only use cv if there are at least 2 folds
-        assert (
-            fold_index is None or self.use_inner_cross_validation is not None
-        ), "fold_index can only be set if cross validation is used"
+        assert fold_index is None or self.use_inner_cross_validation is not None, (
+            "fold_index can only be set if cross validation is used"
+        )
         if fold_index is not None and self.inner_k_folds is not None:
-            assert (
-                fold_index < self.inner_k_folds
-            ), "fold_index can't be larger than the total number of folds"
+            assert fold_index < self.inner_k_folds, (
+                "fold_index can't be larger than the total number of folds"
+            )
         self.fold_index = fold_index
         self._base_dir = base_dir
         self.n_token_limit = n_token_limit
@@ -137,9 +137,9 @@ def num_of_labels(self):
 
     @property
     def feature_vector_size(self):
-        assert (
-            self._feature_vector_size is not None
-        ), "size of feature vector must be set"
+        assert self._feature_vector_size is not None, (
+            "size of feature vector must be set"
+        )
         return self._feature_vector_size
 
     @property
@@ -1252,7 +1252,20 @@ def load_processed_data(
         # If filename is provided
         return self.load_processed_data_from_file(filename)
 
-    def load_processed_data_from_file(self, filename):
+    def load_processed_data_from_file(self, filename: str) -> list[dict[str, Any]]:
+        """Load processed data from a file.
+
+        The full path is not required; only the filename is needed, as it will be joined with the processed directory.
+
+        Args:
+            filename (str): The name of the file to load the processed data from.
+
+        Returns:
+            List[Dict[str, Any]]: The loaded processed data.
+
+        Example:
+            data = self.load_processed_data_from_file('data.pt')
+        """
         return torch.load(
             os.path.join(self.processed_dir, filename), weights_only=False
         )