braindecode 1.3.0.dev181065563__py3-none-any.whl → 1.3.0.dev181594385__py3-none-any.whl

braindecode/models/signal_jepa.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 
 import math
 from copy import deepcopy
-from typing import Any, Sequence
+from pathlib import Path
+from typing import Any, Optional, Sequence
 
 import torch
 from einops.layers.torch import Rearrange
@@ -145,6 +146,8 @@ class _BaseSignalJEPA(EEGModuleMixin, nn.Module):
 class SignalJEPA(_BaseSignalJEPA):
     """Architecture introduced in signal-JEPA for self-supervised pre-training, Guetschel, P et al (2024) [1]_
 
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
+
     This model is not meant for classification but for SSL pre-training.
     Its output shape depends on the input shape.
     For classification purposes, three variants of this model are available:
@@ -231,6 +234,8 @@ class SignalJEPA(_BaseSignalJEPA):
 class SignalJEPA_Contextual(_BaseSignalJEPA):
     """Contextual downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
 
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
+
     This architecture is one of the variants of :class:`SignalJEPA`
     that can be used for classification purposes.
 
@@ -319,25 +324,50 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
     @classmethod
     def from_pretrained(
         cls,
-        model: SignalJEPA,
-        n_outputs: int,
+        model: Optional[SignalJEPA | str | Path] = None,  # type: ignore
+        n_outputs: Optional[int] = None,  # type: ignore
         n_spat_filters: int = 4,
-        chs_info: list[dict[str, Any]] | None = None,
+        chs_info: Optional[list[dict[str, Any]]] = None,  # type: ignore
+        **kwargs,
     ):
-        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model.
+        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model or from Hub.
 
         Parameters
         ----------
-        model: SignalJEPA
-            Pre-trained model.
-        n_outputs: int
-            Number of classes for the new model.
+        model: SignalJEPA, str, Path, or None
+            Either a pre-trained :class:`SignalJEPA` model, a string/Path to a local directory
+            (for Hub-style loading), or None (for Hub loading via kwargs).
+        n_outputs: int or None
+            Number of classes for the new model. Required when loading from a SignalJEPA model,
+            optional when loading from Hub (will be read from config).
         n_spat_filters: int
             Number of spatial filters.
         chs_info: list of dict | None
             Information about each individual EEG channel. This should be filled with
             ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+        **kwargs
+            Additional keyword arguments passed to the parent class for Hub loading.
         """
+        # Check if this is a Hub-style load (from a directory path)
+        if isinstance(model, (str, Path)) or (model is None and kwargs):
+            # This is a Hub load, delegate to parent class
+            if isinstance(model, (str, Path)):
+                # model is actually the repo_id or directory path
+                return super().from_pretrained(model, **kwargs)
+            else:
+                # model is None, treat as hub-style load
+                return super().from_pretrained(**kwargs)
+
+        # This is the original SignalJEPA transfer learning case
+        if not isinstance(model, SignalJEPA):
+            raise TypeError(
+                f"model must be a SignalJEPA instance, a path string, or Path object, got {type(model)}"
+            )
+        if n_outputs is None:
+            raise ValueError(
+                "n_outputs must be provided when loading from a SignalJEPA model"
+            )
+
         feature_encoder = model.feature_encoder
         pos_encoder = model.pos_encoder
         transformer = model.transformer
@@ -377,6 +407,8 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
 class SignalJEPA_PostLocal(_BaseSignalJEPA):
     """Post-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
 
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
+
     This architecture is one of the variants of :class:`SignalJEPA`
     that can be used for classification purposes.
 
@@ -463,22 +495,47 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
 
     @classmethod
     def from_pretrained(
-        cls, model: SignalJEPA, n_outputs: int, n_spat_filters: int = 4
+        cls,
+        model: SignalJEPA | str | Path = None,  # type: ignore
+        n_outputs: int = None,  # type: ignore
+        n_spat_filters: int = 4,
+        **kwargs,
     ):
-        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model.
+        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model or from Hub.
 
         Parameters
         ----------
-        model: SignalJEPA
-            Pre-trained model.
-        n_outputs: int
-            Number of classes for the new model.
+        model: SignalJEPA, str, Path, or None
+            Either a pre-trained :class:`SignalJEPA` model, a string/Path to a local directory
+            (for Hub-style loading), or None (for Hub loading via kwargs).
+        n_outputs: int or None
+            Number of classes for the new model. Required when loading from a SignalJEPA model,
+            optional when loading from Hub (will be read from config).
         n_spat_filters: int
             Number of spatial filters.
-        chs_info: list of dict | None
-            Information about each individual EEG channel. This should be filled with
-            ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+        **kwargs
+            Additional keyword arguments passed to the parent class for Hub loading.
         """
+        # Check if this is a Hub-style load (from a directory path)
+        if isinstance(model, (str, Path)) or (model is None and kwargs):
+            # This is a Hub load, delegate to parent class
+            if isinstance(model, (str, Path)):
+                # model is actually the repo_id or directory path
+                return super().from_pretrained(model, **kwargs)
+            else:
+                # model is None, treat as hub-style load
+                return super().from_pretrained(**kwargs)
+
+        # This is the original SignalJEPA transfer learning case
+        if not isinstance(model, SignalJEPA):
+            raise TypeError(
+                f"model must be a SignalJEPA instance, a path string, or Path object, got {type(model)}"
+            )
+        if n_outputs is None:
+            raise ValueError(
+                "n_outputs must be provided when loading from a SignalJEPA model"
+            )
+
         feature_encoder = model.feature_encoder
         assert feature_encoder is not None
         new_model = cls(
@@ -501,6 +558,8 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
 class SignalJEPA_PreLocal(_BaseSignalJEPA):
     """Pre-local downstream architecture introduced in signal-JEPA Guetschel, P et al (2024) [1]_.
 
+    :bdg-success:`Convolution` :bdg-dark-line:`Channel` :bdg-danger:`Large Brain Model`
+
     This architecture is one of the variants of :class:`SignalJEPA`
     that can be used for classification purposes.
 
@@ -597,22 +656,47 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
 
     @classmethod
     def from_pretrained(
-        cls, model: SignalJEPA, n_outputs: int, n_spat_filters: int = 4
+        cls,
+        model: SignalJEPA | str | Path = None,  # type: ignore
+        n_outputs: int = None,  # type: ignore
+        n_spat_filters: int = 4,
+        **kwargs,
     ):
-        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model.
+        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model or from Hub.
 
         Parameters
         ----------
-        model: SignalJEPA
-            Pre-trained model.
-        n_outputs: int
-            Number of classes for the new model.
+        model: SignalJEPA, str, Path, or None
+            Either a pre-trained :class:`SignalJEPA` model, a string/Path to a local directory
+            (for Hub-style loading), or None (for Hub loading via kwargs).
+        n_outputs: int or None
+            Number of classes for the new model. Required when loading from a SignalJEPA model,
+            optional when loading from Hub (will be read from config).
         n_spat_filters: int
             Number of spatial filters.
-        chs_info: list of dict | None
-            Information about each individual EEG channel. This should be filled with
-            ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+        **kwargs
+            Additional keyword arguments passed to the parent class for Hub loading.
         """
+        # Check if this is a Hub-style load (from a directory path)
+        if isinstance(model, (str, Path)) or (model is None and kwargs):
+            # This is a Hub load, delegate to parent class
+            if isinstance(model, (str, Path)):
+                # model is actually the repo_id or directory path
+                return super().from_pretrained(model, **kwargs)
+            else:
+                # model is None, treat as hub-style load
+                return super().from_pretrained(**kwargs)
+
+        # This is the original SignalJEPA transfer learning case
+        if not isinstance(model, SignalJEPA):
+            raise TypeError(
+                f"model must be a SignalJEPA instance, a path string, or Path object, got {type(model)}"
+            )
+        if n_outputs is None:
+            raise ValueError(
+                "n_outputs must be provided when loading from a SignalJEPA model"
+            )
+
         feature_encoder = model.feature_encoder
         assert feature_encoder is not None
         new_model = cls(

braindecode/models/sinc_shallow.py

@@ -12,6 +12,8 @@ from braindecode.models.base import EEGModuleMixin
 class SincShallowNet(EEGModuleMixin, nn.Module):
     """Sinc-ShallowNet from Borra, D et al (2020) [borra2020]_.
 
+    :bdg-success:`Convolution` :bdg-warning:`Interpretability`
+
     .. figure:: https://ars.els-cdn.com/content/image/1-s2.0-S0893608020302021-gr2_lrg.jpg
         :align: center
         :alt: SincShallowNet Architecture
@@ -19,23 +21,24 @@ class SincShallowNet(EEGModuleMixin, nn.Module):
     The Sinc-ShallowNet architecture has these fundamental blocks:
 
     1. **Block 1: Spectral and Spatial Feature Extraction**
-        - *Temporal Sinc-Convolutional Layer*:
-            Uses parametrized sinc functions to learn band-pass filters,
-            significantly reducing the number of trainable parameters by only
-            learning the lower and upper cutoff frequencies for each filter.
-       - *Spatial Depthwise Convolutional Layer*:
-            Applies depthwise convolutions to learn spatial filters for
-            each temporal feature map independently, further reducing
-            parameters and enhancing interpretability.
-       - *Batch Normalization*
+
+        - *Temporal Sinc-Convolutional Layer*: Uses parametrized sinc functions to learn band-pass filters,
+          significantly reducing the number of trainable parameters by only
+          learning the lower and upper cutoff frequencies for each filter.
+        - *Spatial Depthwise Convolutional Layer*: Applies depthwise convolutions to learn spatial filters for
+          each temporal feature map independently, further reducing
+          parameters and enhancing interpretability.
+        - *Batch Normalization*
 
     2. **Block 2: Temporal Aggregation**
+
         - *Activation Function*: ELU
         - *Average Pooling Layer*: Aggregation by averaging spatial dim
         - *Dropout Layer*
         - *Flatten Layer*
 
     3. **Block 3: Classification**
+
         - *Fully Connected Layer*: Maps the feature vector to n_outputs.
 
     **Implementation Notes:**