braindecode 0.8__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/training/scoring.py

@@ -0,0 +1,483 @@
+# Authors: Maciej Sliwowski <maciek.sliwowski@gmail.com>
+#          Robin Tibor Schirrmeister <robintibor@gmail.com>
+#          Alexandre Gramfort <alexandre.gramfort@inria.fr>
+#          Lukas Gemein <l.gemein@gmail.com>
+#          Mohammed Fattouh <mo.fattouh@gmail.com>
+#
+# License: BSD-3
+
+import warnings
+from contextlib import contextmanager
+
+import numpy as np
+import torch
+from mne.utils.check import check_version
+from skorch.callbacks.scoring import EpochScoring
+from skorch.dataset import unpack_data
+from skorch.utils import to_numpy
+from torch.utils.data import DataLoader
+
+
+def trial_preds_from_window_preds(preds, i_window_in_trials, i_stop_in_trials):
+    """
+    Assigning window predictions to trials  while removing duplicate
+    predictions.
+
+    Parameters
+    ----------
+    preds: list of ndarrays (at least 2darrays)
+        List of window predictions, in each window prediction
+         time is in axis=1
+    i_window_in_trials: list
+        Index/number of window in trial
+    i_stop_in_trials: list
+        stop position of window in trial
+
+    Returns
+    -------
+    preds_per_trial: list of ndarrays
+        Predictions in each trial, duplicates removed
+
+    """
+    assert len(preds) == len(i_window_in_trials) == len(i_stop_in_trials), (
+        f"{len(preds)}, {len(i_window_in_trials)}, {len(i_stop_in_trials)}"
+    )
+
+    # Algorithm for assigning window predictions to trials
+    # while removing duplicate predictions:
+    # Loop through windows:
+    # In each iteration you have predictions (assumed: #classes x #timesteps,
+    # or at least #timesteps must be in axis=1)
+    # and you have i_window_in_trial, i_stop_in_trial
+    # (i_trial removed from variable names for brevity)
+    # You first check if the i_window_in_trial is 1 larger
+    # than in last iteration, then you are still in the same trial
+    # Otherwise you are in a new trial
+    # If you are in the same trial, you check for duplicate predictions
+    # Only take predictions that are after (inclusive)
+    # the stop of the last iteration (i.e., the index of final prediction
+    # in the last iteration)
+    # Then add the duplicate-removed predictions from this window
+    # to predictions for current trial
+    preds_per_trial = []
+    cur_trial_preds = []
+    i_last_stop = None
+    i_last_window = -1
+    for window_preds, i_window, i_stop in zip(
+        preds, i_window_in_trials, i_stop_in_trials
+    ):
+        window_preds = np.array(window_preds)
+        if i_window != (i_last_window + 1):
+            assert i_window == 0, "window numbers in new trial should start from 0"
+            preds_per_trial.append(np.concatenate(cur_trial_preds, axis=1))
+            cur_trial_preds = []
+            i_last_stop = None
+
+        if i_last_stop is not None:
+            # Remove duplicates
+            n_needed_preds = i_stop - i_last_stop
+            window_preds = window_preds[:, -n_needed_preds:]
+        cur_trial_preds.append(window_preds)
+        i_last_window = i_window
+        i_last_stop = i_stop
+    # add last trial preds
+    preds_per_trial.append(np.concatenate(cur_trial_preds, axis=1))
+    return preds_per_trial
+
+
+@contextmanager
+def _cache_net_forward_iter(net, use_caching, y_preds):
+    """Caching context for ``skorch.NeuralNet`` instance.
+    Returns a modified version of the net whose ``forward_iter``
+    method will subsequently return cached predictions. Leaving the
+    context will undo the overwrite of the ``forward_iter`` method.
+    """
+    if not use_caching:
+        yield net
+        return
+    y_preds = iter(y_preds)
+
+    # pylint: disable=unused-argument
+    def cached_forward_iter(*args, device=net.device, **kwargs):
+        for yp in y_preds:
+            yield yp.to(device=device)
+
+    net.forward_iter = cached_forward_iter
+    try:
+        yield net
+    finally:
+        # By setting net.forward_iter we define an attribute
+        # `forward_iter` that precedes the bound method
+        # `forward_iter`. By deleting the entry from the attribute
+        # dict we undo this.
+        del net.__dict__["forward_iter"]
+
+
+class CroppedTrialEpochScoring(EpochScoring):
+    """
+    Class to compute scores for trials from a model that predicts (super)crops.
+    """
+
+    # XXX needs a docstring !!!
+
+    def __init__(
+        self,
+        scoring,
+        lower_is_better=True,
+        on_train=False,
+        name=None,
+        target_extractor=to_numpy,
+        use_caching=True,
+    ):
+        super().__init__(
+            scoring=scoring,
+            lower_is_better=lower_is_better,
+            on_train=on_train,
+            name=name,
+            target_extractor=target_extractor,
+            use_caching=use_caching,
+        )
+        if not self.on_train:
+            self.window_inds_ = []
+
+    def _initialize_cache(self):
+        super()._initialize_cache()
+        self.crops_to_trials_computed = False
+        self.y_trues_ = []
+        self.y_preds_ = []
+        if not self.on_train:
+            self.window_inds_ = []
+
+    def on_batch_end(self, net, batch, y_pred, training, **kwargs):
+        # Skorch saves the predictions without moving them from GPU
+        # https://github.com/skorch-dev/skorch/blob/fe71e3d55a4ae5f5f94ef7bdfc00fca3b3fd267f/skorch/callbacks/scoring.py#L385
+        # This can cause memory issues in case of a large number of predictions
+        # Therefore here we move them to CPU already
+        super().on_batch_end(net, batch, y_pred, training, **kwargs)
+        if self.use_caching and training == self.on_train:
+            self.y_preds_[-1] = self.y_preds_[-1].cpu()
+
+    def on_epoch_end(self, net, dataset_train, dataset_valid, **kwargs):
+        assert self.use_caching
+        if not self.crops_to_trials_computed:
+            if self.on_train:
+                # Prevent that rng state of torch is changed by
+                # creation+usage of iterator
+                rng_state = torch.random.get_rng_state()
+                pred_results = net.predict_with_window_inds_and_ys(dataset_train)
+                torch.random.set_rng_state(rng_state)
+            else:
+                pred_results = {}
+                pred_results["i_window_in_trials"] = np.concatenate(
+                    [i[0].cpu().numpy() for i in self.window_inds_]
+                )
+                pred_results["i_window_stops"] = np.concatenate(
+                    [i[2].cpu().numpy() for i in self.window_inds_]
+                )
+                pred_results["preds"] = np.concatenate(
+                    [y_pred.cpu().numpy() for y_pred in self.y_preds_]
+                )
+                pred_results["window_ys"] = np.concatenate(
+                    [y.cpu().numpy() for y in self.y_trues_]
+                )
+
+            # A new trial starts
+            # when the index of the window in trials
+            # does not increment by 1
+            # Add dummy infinity at start
+            window_0_per_trial_mask = (
+                np.diff(pred_results["i_window_in_trials"], prepend=[np.inf]) != 1
+            )
+            trial_ys = pred_results["window_ys"][window_0_per_trial_mask]
+            trial_preds = trial_preds_from_window_preds(
+                pred_results["preds"],
+                pred_results["i_window_in_trials"],
+                pred_results["i_window_stops"],
+            )
+
+            # Average across the timesteps of each trial so we have per-trial
+            # predictions already, these will be just passed through the forward
+            # method of the classifier/regressor to the skorch scoring function.
+            # trial_preds is a list, each item is a 2d array classes x time
+            y_preds_per_trial = np.array([np.mean(p, axis=1) for p in trial_preds])
+            # Move into format expected by skorch (list of torch tensors)
+            y_preds_per_trial = [torch.tensor(y_preds_per_trial)]
+
+            # Store the computed trial preds for all Cropped Callbacks
+            # that are also on same set
+            cbs = net.callbacks_
+            epoch_cbs = [
+                cb
+                for name, cb in cbs
+                if isinstance(cb, CroppedTrialEpochScoring)
+                and (cb.on_train == self.on_train)
+            ]
+            for cb in epoch_cbs:
+                cb.y_preds_ = y_preds_per_trial
+                cb.y_trues_ = trial_ys
+                cb.crops_to_trials_computed = True
+
+        dataset = dataset_train if self.on_train else dataset_valid
+
+        with _cache_net_forward_iter(
+            net, self.use_caching, self.y_preds_
+        ) as cached_net:
+            current_score = self._scoring(cached_net, dataset, self.y_trues_)
+        self._record_score(net.history, current_score)
+
+        return
+
+
+class CroppedTimeSeriesEpochScoring(CroppedTrialEpochScoring):
+    """
+    Class to compute scores for trials from a model that predicts (super)crops with
+    time series target.
+    """
+
+    def on_epoch_end(self, net, dataset_train, dataset_valid, **kwargs):
+        assert self.use_caching
+        if not self.crops_to_trials_computed:
+            if self.on_train:
+                # Prevent that rng state of torch is changed by
+                # creation+usage of iterator
+                rng_state = torch.random.get_rng_state()
+                pred_results = net.predict_with_window_inds_and_ys(dataset_train)
+                torch.random.set_rng_state(rng_state)
+            else:
+                pred_results = {}
+                pred_results["i_window_in_trials"] = np.concatenate(
+                    [i[0].cpu().numpy() for i in self.window_inds_]
+                )
+                pred_results["i_window_stops"] = np.concatenate(
+                    [i[2].cpu().numpy() for i in self.window_inds_]
+                )
+                pred_results["preds"] = np.concatenate(
+                    [y_pred.cpu().numpy() for y_pred in self.y_preds_]
+                )
+                pred_results["window_ys"] = np.concatenate(
+                    [y.cpu().numpy() for y in self.y_trues_]
+                )
+
+            num_preds = pred_results["preds"][-1].shape[-1]
+            # slice the targets to fit preds shape
+            pred_results["window_ys"] = [
+                targets[:, -num_preds:] for targets in pred_results["window_ys"]
+            ]
+
+            trial_preds = trial_preds_from_window_preds(
+                pred_results["preds"],
+                pred_results["i_window_in_trials"],
+                pred_results["i_window_stops"],
+            )
+
+            trial_ys = trial_preds_from_window_preds(
+                pred_results["window_ys"],
+                pred_results["i_window_in_trials"],
+                pred_results["i_window_stops"],
+            )
+
+            # the output is a list of predictions/targets per trial where each item is a
+            # timeseries of predictions/targets of shape (n_classes x timesteps)
+
+            # mask NaNs form targets
+            preds = np.hstack(trial_preds)  # n_classes x timesteps in all trials
+            targets = np.hstack(trial_ys)
+            # create valid targets mask
+            mask = ~np.isnan(targets)
+            # select valid targets that have a matching predictions
+            masked_targets = targets[mask]
+            # For classification there is only one row in targets and n_classes rows in preds
+            if mask.shape[0] != preds.shape[0]:
+                masked_preds = preds[:, mask[0, :]]
+            else:
+                masked_preds = preds[mask]
+
+            # Store the computed trial preds for all Cropped Callbacks
+            # that are also on same set
+            cbs = net.callbacks_
+            epoch_cbs = [
+                cb
+                for name, cb in cbs
+                if isinstance(cb, CroppedTimeSeriesEpochScoring)
+                and (cb.on_train == self.on_train)
+            ]
+            masked_preds = [torch.tensor(masked_preds.T)]
+            for cb in epoch_cbs:
+                cb.y_preds_ = masked_preds
+                cb.y_trues_ = masked_targets.T
+                cb.crops_to_trials_computed = True
+
+        dataset = dataset_train if self.on_train else dataset_valid
+
+        with _cache_net_forward_iter(
+            net, self.use_caching, self.y_preds_
+        ) as cached_net:
+            current_score = self._scoring(cached_net, dataset, self.y_trues_)
+        self._record_score(net.history, current_score)
+
+
+class PostEpochTrainScoring(EpochScoring):
+    """
+    Epoch Scoring class that recomputes predictions after the epoch
+    on the training in validation mode.
+
+    Note: For unknown reasons, this affects global random generator and
+    therefore all results may change slightly if you add this scoring callback.
+
+    Parameters
+    ----------
+    scoring : None, str, or callable (default=None)
+      If None, use the ``score`` method of the model. If str, it
+      should be a valid sklearn scorer (e.g. "f1", "accuracy"). If a
+      callable, it should have the signature (model, X, y), and it
+      should return a scalar. This works analogously to the
+      ``scoring`` parameter in sklearn's ``GridSearchCV`` et al.
+    lower_is_better : bool (default=True)
+      Whether lower scores should be considered better or worse.
+    name : str or None (default=None)
+      If not an explicit string, tries to infer the name from the
+      ``scoring`` argument.
+    target_extractor : callable (default=to_numpy)
+      This is called on y before it is passed to scoring.
+    """
+
+    def __init__(
+        self,
+        scoring,
+        lower_is_better=True,
+        name=None,
+        target_extractor=to_numpy,
+    ):
+        super().__init__(
+            scoring=scoring,
+            lower_is_better=lower_is_better,
+            on_train=True,
+            name=name,
+            target_extractor=target_extractor,
+            use_caching=False,
+        )
+
+    def on_epoch_end(self, net, dataset_train, dataset_valid, **kwargs):
+        if len(self.y_preds_) == 0:
+            dataset = net.get_dataset(dataset_train)
+            # Prevent that rng state of torch is changed by
+            # creation+usage of iterator
+            # Unfortunatenly calling __iter__() of a pytorch
+            # DataLoader will change the random state
+            # Note line below setting rng state back
+            rng_state = torch.random.get_rng_state()
+            iterator = net.get_iterator(dataset, training=False)
+            y_preds = []
+            y_test = []
+            for batch in iterator:
+                batch_X, batch_y = unpack_data(batch)
+                # TODO: remove after skorch 0.10 release
+                if not check_version("skorch", min_version="0.10.1"):
+                    yp = net.evaluation_step(batch_X, training=False)
+                # X, y unpacking has been pushed downstream in skorch 0.10
+                else:
+                    yp = net.evaluation_step(batch, training=False)
+                yp = yp.to(device="cpu")
+                y_test.append(self.target_extractor(batch_y))
+                y_preds.append(yp)
+            y_test = np.concatenate(y_test)
+            torch.random.set_rng_state(rng_state)
+
+            # Adding the recomputed preds to all other
+            # instances of PostEpochTrainScoring of this
+            # Skorch-Net (NeuralNet, BraindecodeClassifier etc.)
+            # (They will be reinitialized to empty lists by skorch
+            # each epoch)
+            cbs = net.callbacks_
+            epoch_cbs = [
+                cb for name, cb in cbs if isinstance(cb, PostEpochTrainScoring)
+            ]
+            for cb in epoch_cbs:
+                cb.y_preds_ = y_preds
+                cb.y_trues_ = y_test
+        # y pred should be same as self.y_preds_
+        # Unclear if this also leads to any
+        # random generator call?
+        with _cache_net_forward_iter(
+            net, use_caching=True, y_preds=self.y_preds_
+        ) as cached_net:
+            current_score = self._scoring(cached_net, dataset_train, self.y_trues_)
+        self._record_score(net.history, current_score)
+
+
+def predict_trials(module, dataset, return_targets=True, batch_size=1, num_workers=0):
+    """Create trialwise predictions and optionally also return trialwise
+    labels from cropped dataset given module.
+
+    Parameters
+    ----------
+    module: torch.nn.Module
+        A pytorch model implementing forward.
+    dataset: braindecode.datasets.BaseConcatDataset
+        A braindecode dataset to be predicted.
+    return_targets: bool
+        If True, additionally returns the trial targets.
+    batch_size: int
+        The batch size used to iterate the dataset.
+    num_workers: int
+        Number of workers used in DataLoader to iterate the dataset.
+
+    Returns
+    -------
+        trial_predictions: np.ndarray
+            3-dimensional array (n_trials x n_classes x n_predictions), where
+            the number of predictions depend on the chosen window size and the
+            receptive field of the network.
+        trial_labels: np.ndarray
+            2-dimensional array (n_trials x n_targets) where the number of
+            targets depends on the decoding paradigm and can be either a single
+            value, multiple values, or a sequence.
+    """
+    # Ensure the model is in evaluation mode
+    module.eval()
+    # we have a cropped dataset if there exists at least one trial with more
+    # than one compute window
+    more_than_one_window = sum(dataset.get_metadata()["i_window_in_trial"] != 0) > 0
+    if not more_than_one_window:
+        warnings.warn(
+            "This function was designed to predict trials from "
+            "cropped datasets, which typically have multiple compute "
+            "windows per trial. The given dataset has exactly one "
+            "window per trial."
+        )
+    loader = DataLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+    )
+    device = next(module.parameters()).device
+    all_preds, all_ys, all_inds = [], [], []
+    with torch.no_grad():
+        for X, y, ind in loader:
+            X = X.to(device)
+            preds = module(X)
+            all_preds.extend(preds.cpu().numpy().astype(np.float32))
+            all_ys.extend(y.cpu().numpy().astype(np.float32))
+            all_inds.extend(ind)
+    preds_per_trial = trial_preds_from_window_preds(
+        preds=all_preds,
+        i_window_in_trials=torch.cat(all_inds[0::3]),
+        i_stop_in_trials=torch.cat(all_inds[2::3]),
+    )
+    preds_per_trial = np.array(preds_per_trial)
+    if return_targets:
+        if all_ys[0].shape == ():
+            all_ys = np.array(all_ys)
+            ys_per_trial = all_ys[
+                np.diff(torch.cat(all_inds[0::3]), prepend=[np.inf]) != 1
+            ]
+        else:
+            ys_per_trial = trial_preds_from_window_preds(
+                preds=all_ys,
+                i_window_in_trials=torch.cat(all_inds[0::3]),
+                i_stop_in_trials=torch.cat(all_inds[2::3]),
+            )
+            ys_per_trial = np.array(ys_per_trial)
+        return preds_per_trial, ys_per_trial
+    return preds_per_trial

braindecode/util.py

@@ -12,6 +12,7 @@ import mne
 import numpy as np
 import torch
 from sklearn.utils import check_random_state
+from torch import Tensor
 
 
 def set_random_seeds(seed, cuda, cudnn_benchmark=None):
@@ -51,7 +52,9 @@ def set_random_seeds(seed, cuda, cudnn_benchmark=None):
                 warn(
                     "torch.backends.cudnn.benchmark was set to True which may results in lack of "
                     "reproducibility. In some cases to ensure reproducibility you may need to "
-                    "set torch.backends.cudnn.benchmark to False.", UserWarning)
+                    "set torch.backends.cudnn.benchmark to False.",
+                    UserWarning,
+                )
         else:
             raise ValueError(
                 f"cudnn_benchmark expected to be bool or None, got '{cudnn_benchmark}'"
@@ -60,19 +63,7 @@ def set_random_seeds(seed, cuda, cudnn_benchmark=None):
     np.random.seed(seed)
 
 
-def np_to_var(
-    X, requires_grad=False, dtype=None, pin_memory=False, **tensor_kwargs
-):
-    warn("np_to_var has been renamed np_to_th, please use np_to_th instead")
-    return np_to_th(
-        X, requires_grad=requires_grad, dtype=dtype, pin_memory=pin_memory,
-        **tensor_kwargs
-    )
-
-
-def np_to_th(
-    X, requires_grad=False, dtype=None, pin_memory=False, **tensor_kwargs
-):
+def np_to_th(X, requires_grad=False, dtype=None, pin_memory=False, **tensor_kwargs):
     """
     Convenience function to transform numpy array to `torch.Tensor`.
 
@@ -103,12 +94,7 @@ def np_to_th(
     return X_tensor
 
 
-def var_to_np(var):
-    warn("var_to_np has been renamed th_to_np, please use th_to_np instead")
-    return th_to_np(var)
-
-
-def th_to_np(var):
+def th_to_np(var: Tensor):
     """Convenience function to transform `torch.Tensor` to numpy
     array.
 
@@ -209,15 +195,11 @@ def wrap_reshape_apply_fn(stat_fn, a, b, axis_a, axis_b):
     )
     assert np.array_equal(n_stat_axis_a, n_stat_axis_b)
     stat_result = stat_fn(flat_topo_a, flat_topo_b)
-    topo_result = stat_result.reshape(
-        tuple(n_other_axis_a) + tuple(n_other_axis_b)
-    )
+    topo_result = stat_result.reshape(tuple(n_other_axis_a) + tuple(n_other_axis_b))
     return topo_result
 
 
-def get_balanced_batches(
-    n_trials, rng, shuffle, n_batches=None, batch_size=None
-):
+def get_balanced_batches(n_trials, rng, shuffle, n_batches=None, batch_size=None):
     """Create indices for batches balanced in size
     (batches will have maximum size difference of 1).
     Supply either batch size or number of batches. Resulting batches
@@ -268,9 +250,17 @@ def get_balanced_batches(
     return batches
 
 
-def create_mne_dummy_raw(n_channels, n_times, sfreq, include_anns=True,
-                         description=None, savedir=None, save_format='fif',
-                         overwrite=True, random_state=None):
+def create_mne_dummy_raw(
+    n_channels,
+    n_times,
+    sfreq,
+    include_anns=True,
+    description=None,
+    savedir=None,
+    save_format="fif",
+    overwrite=True,
+    random_state=None,
+):
     """Create an mne.io.RawArray with fake data, and optionally save it.
 
     This will overwrite already existing files.
@@ -305,20 +295,21 @@ def create_mne_dummy_raw(n_channels, n_times, sfreq, include_anns=True,
     """
     random_state = check_random_state(random_state)
     data = random_state.rand(n_channels, n_times)
-    ch_names = [f'ch{i}' for i in range(n_channels)]
-    ch_types = ['eeg'] * n_channels
+    ch_names = [f"ch{i}" for i in range(n_channels)]
+    ch_types = ["eeg"] * n_channels
     info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
 
     raw = mne.io.RawArray(data, info)
 
     if include_anns:
         n_anns = 10
-        inds = np.linspace(
-            int(sfreq * 2), int(n_times - sfreq * 2), num=n_anns).astype(int)
+        inds = np.linspace(int(sfreq * 2), int(n_times - sfreq * 2), num=n_anns).astype(
+            int
+        )
         onset = raw.times[inds]
         duration = [1] * n_anns
         if description is None:
-            description = ['test'] * n_anns
+            description = ["test"] * n_anns
         anns = mne.Annotations(onset, duration, description)
         raw = raw.set_annotations(anns)
 
@@ -326,18 +317,17 @@ def create_mne_dummy_raw(n_channels, n_times, sfreq, include_anns=True,
     if savedir is not None:
         if not isinstance(save_format, list):
             save_format = [save_format]
-        fname = os.path.join(savedir, 'fake_eeg_raw')
+        fname = os.path.join(savedir, "fake_eeg_raw")
 
-        if 'fif' in save_format:
-            fif_fname = fname + '.fif'
+        if "fif" in save_format:
+            fif_fname = fname + ".fif"
             raw.save(fif_fname, overwrite=overwrite)
-            save_fname['fif'] = fif_fname
-        if 'hdf5' in save_format:
-            h5_fname = fname + '.h5'
-            with h5py.File(h5_fname, 'w') as f:
-                f.create_dataset(
-                    'fake_raw', dtype='f8', data=raw.get_data())
-            save_fname['hdf5'] = h5_fname
+            save_fname["fif"] = fif_fname
+        if "hdf5" in save_format:
+            h5_fname = fname + ".h5"
+            with h5py.File(h5_fname, "w") as f:
+                f.create_dataset("fake_raw", dtype="f8", data=raw.get_data())
+            save_fname["hdf5"] = h5_fname
 
     return raw, save_fname
 
@@ -349,7 +339,9 @@ class ThrowAwayIndexLoader(object):
         self.last_i = None
         self.is_regression = is_regression
 
-    def __iter__(self, ):
+    def __iter__(
+        self,
+    ):
         normal_iter = self.loader.__iter__()
         for batch in normal_iter:
             if len(batch) == 3:
@@ -360,7 +352,7 @@ class ThrowAwayIndexLoader(object):
                 x, y = batch
 
             # TODO: should be on dataset side
-            if hasattr(x, 'type'):
+            if hasattr(x, "type"):
                 x = x.type(torch.float32)
                 if self.is_regression:
                     y = y.type(torch.float32)
@@ -370,23 +362,26 @@ class ThrowAwayIndexLoader(object):
 
 
 def update_estimator_docstring(base_class, docstring):
-    base_doc = base_class.__doc__.replace(' : ', ': ')
-    idx = base_doc.find('callbacks:')
-    idx_end = idx + base_doc[idx:].find('\n\n')
+    base_doc = base_class.__doc__.replace(" : ", ": ")
+    idx = base_doc.find("callbacks:")
+    idx_end = idx + base_doc[idx:].find("\n\n")
     # remove callback descripiton already included in braindecode docstring
-    filtered_doc = base_doc[:idx] + base_doc[idx_end + 6:]
-    splitted = docstring.split('Parameters\n    ----------\n    ')
+    filtered_doc = base_doc[:idx] + base_doc[idx_end + 6 :]
+    splitted = docstring.split("Parameters\n    ----------\n    ")
     out_docstring = (
-        splitted[0] +
-        filtered_doc[filtered_doc.find('Parameters'):filtered_doc.find('Attributes')] +
-        splitted[1] +
-        filtered_doc[filtered_doc.find('Attributes'):])
+        splitted[0]
+        + filtered_doc[
+            filtered_doc.find("Parameters") : filtered_doc.find("Attributes")
+        ]
+        + splitted[1]
+        + filtered_doc[filtered_doc.find("Attributes") :]
+    )
     return out_docstring
 
 
 def _update_moabb_docstring(base_class, docstring):
     base_doc = base_class.__doc__
-    out_docstring = base_doc + f'\n\n{docstring}'
+    out_docstring = base_doc + f"\n\n{docstring}"
     return out_docstring
 
 
@@ -406,8 +401,9 @@ def read_all_file_names(directory, extension):
     file_paths: list(str)
         List of all files found in (sub)directories of path.
     """
-    assert extension.startswith('.')
-    file_paths = glob.glob(directory + '**/*' + extension, recursive=True)
+    assert extension.startswith(".")
+    file_paths = glob.glob(directory + "**/*" + extension, recursive=True)
     assert len(file_paths) > 0, (
-        f'something went wrong. Found no {extension} files in {directory}')
+        f"something went wrong. Found no {extension} files in {directory}"
+    )
     return file_paths