Fix reset_head leaving the saved config on the previous head - #1181

Master's braindecode#1181 requires reset_head implementations to record the new
n_outputs with _set_n_outputs, so get_config() and save_pretrained()
describe the head the model has. SleepFM and SleepFMStager still set
_n_outputs directly, which the merged test_reset_head_updates_config and
test_reset_head_model_reloads_after_saving cases catch.

* feat(models): add SleepFM model for multimodal polysomnography

Introduce the SleepFM model, which learns channel-agnostic representations from polysomnography data, including EEG, EOG, respiratory signals, ECG, and EMG. The model utilizes a convolutional tokenizer to process non-overlapping patches of input signals and incorporates attention mechanisms for pooling variable channel sets. Update summary.csv to include SleepFM and its parameters.

* feat(models): include SleepFM and SleepFMStager in model exports

Add SleepFM and SleepFMStager to the braindecode.models module, updating the __all__ export to ensure they are accessible for users. This change enhances the model offerings for multimodal polysomnography analysis.

* feat(models): add SleepFM and SleepFMStager parameters to model exports

Enhance the braindecode.models module by including mandatory parameters for the SleepFM and SleepFMStager models. This update ensures that both models are fully integrated into the existing framework, facilitating their use in multimodal polysomnography analysis.

* docs: update API and what's new for SleepFM and SleepFMStager

Enhance documentation by adding details about the new SleepFM and SleepFMStager models, including their functionalities and compatibility with existing checkpoints. Update the models categorization to reflect their inclusion and provide a comprehensive overview in the what's new section.

* feat(tests): add unit tests for SleepFM and SleepFMStager models

Introduce comprehensive unit tests for the SleepFM and SleepFMStager models, validating their functionality, including tokenizer behavior, attention pooling, and model output shapes. Ensure that the models correctly handle various input scenarios and maintain expected output dimensions, enhancing the test coverage for these new additions to the braindecode framework.

* docs: update NOTICE.txt to include SleepFM licensing information

Add licensing details for the SleepFM model, indicating its adaptation from `zou-group/sleepfm-clinical` under CC BY-NC 4.0. This update ensures proper attribution and compliance with licensing terms for the newly integrated model.

* refactor(models): enhance mask validation in _SleepFMAttentionPooling

Update the _validate_mask method to accept a mask_name parameter for improved error messaging. This change allows for more descriptive exceptions when validating key padding and channel masks, enhancing code clarity and maintainability.

* feat(tests): add additional unit tests for SleepFM attention pooling and channel mask validation

Introduce a new test for compiled attention pooling to ensure it correctly rejects all masked samples. Additionally, add parameterized tests to validate error handling for various channel mask shapes, enhancing the robustness of the SleepFM model's input validation.

* fix(models): enhance handling of all masked samples in _SleepFMAttentionPooling

Update the _SleepFMAttentionPooling class to safely handle cases where all input channels are masked. The changes ensure that outputs are correctly masked to zero when all channels are invalid, improving the robustness of the attention pooling mechanism. Additionally, modify the corresponding unit test to verify this behavior, ensuring that the model behaves as expected under these conditions.

* fix(tests): update SleepFM checkpoint URLs to Hugging Face

Modify unit tests for SleepFM to replace outdated GitHub URLs with updated links from Hugging Face. This change ensures that the tests reference the correct model checkpoints, improving the reliability of the test suite.

* docs(models): clarify SleepFMStager documentation with reference to downstream architecture

* FIX preserve SleepFM singleton parity

* Simplify SleepFM and load its weights through from_pretrained

Answers the review on braindecode/models/sleepfm.py:

- Public models first, private helper classes after, per the convention of
  the recently merged ports (mvpformer.py, steegformer.py).
- The two mask helpers become module-level functions: the stager was reaching
  into `SleepFM._prepare_channel_mask` and the pooling into
  `_SleepFMAttentionPooling._validate_mask`, neither of which needs a class.
  As a side effect the validation errors now quote `channel_mask`, the public
  argument, instead of the internal `key_padding_mask`.
- The tokenizer owns the patch arithmetic. Both constructors asked
  `n_times // patch_size` themselves and re-implemented the "at least one
  complete patch" check; they now ask `patch_embedding.n_patches()`.
- A wrong `sfreq` warns instead of raising. 128 Hz is what the released
  weights were trained at, not a constraint of the architecture, so the model
  stays usable on other data at the user's risk.
- einops for every channel/patch fold, which is where the shapes were hardest
  to follow, plus the math behind `encode` and the three stages of the
  staging head in the docstrings.
- `tokenize()` was a one-line delegation to `patch_embedding` and is gone.
- The bespoke checkpoint API (`load_pretrained_backbone`,
  `load_pretrained_staging_head` and the four key-remapping helpers) is
  replaced by `from_pretrained`, as for ZUNA and LUNA. The keys are rewritten
  once when mirroring the weights, so the library carries no remapping code,
  and the stager mirror merges the released tokenizer with the released
  staging head -- one call now returns a fully pretrained stager instead of
  two.
- `# nosec B105`: Bandit reads any dict key containing "token" with a
  constant value as a hardcoded credential, and `{"features": ...,
  "cls_token": None}` is braindecode's feature-return contract, shared with
  eight other models.

The refactor is numerically neutral: same state_dict and same input give
bit-exact outputs before and after, logits, features and masked forward alike.

* Fold the SleepFM tests into the shared model suites

`test_sleepfm.py` is deleted: shapes, the feature contract, `reset_head`,
compilation, export and registration are already covered for every model by
the parametrized suites, and duplicating them per model is what the shared
files exist to avoid.

What survives moves to test_foundation_models.py, next to the other foundation
models, and is only what is specific to SleepFM: the channel mask. A PSG
montage varies between subjects, so the tests worth keeping are that a masked
channel cannot reach any output whatever it contains, that a fully masked
sample is refused, and that the errors name `channel_mask`.

SleepFM and SleepFMStager also leave the `test_model_compiled` exclusion list.
The stated reason -- official-size configurations being too expensive to
compile -- does not hold: both compile and run together in 23 s on CPU at the
summary.csv size. They stay excluded from `test_model_torch_script`, with the
reason corrected to the measured one: TorchScript cannot type the
`tuple(x.shape)` shape-guard messages, and rejects the polymorphic
Dict/Tensor return exactly as it does for EEGDINO and MVPFormer.

The two network tests now go through `from_pretrained`, and additionally check
that the stager's tokenizer, recurrent head and output layer all come from the
checkpoint rather than from the random initialisation.

* Trim the SleepFM changelog entry

The entry described the loading API that no longer exists and the 128 Hz
check that is now a warning. It also spent five lines on the paper's
disease-prediction baselines; one sentence saying why they are out of scope
(they need age, sex, BMI and race/ethnicity, which are not
electrophysiological inputs) says the same thing.

* FIX: SleepFM reset_head keeps the saved config in sync

Master's #1181 requires reset_head implementations to record the new
n_outputs with _set_n_outputs, so get_config() and save_pretrained()
describe the head the model has. SleepFM and SleepFMStager still set
_n_outputs directly, which the merged test_reset_head_updates_config and
test_reset_head_model_reloads_after_saving cases catch.

* FIX: SleepFMStager reproduces the released staging pipeline

SleepFMStager.forward fed raw tokenizer outputs to the staging head,
whereas the released staging head was trained on SleepFM encoder
embeddings computed per modality on 5-minute chunks; its logits differed
from the official pipeline by up to 3.18 on released weights. It also had
no temporal mask, and in training mode masked channels entered the
tokenizer's BatchNorm statistics.

- The stager now holds the SleepFM encoder (same parameter names, minus
  the trial-level temporal pooling). Channels are grouped with the new
  channel_modalities argument; each modality is encoded chunk by chunk
  (encoder_chunk_patches=60), positions restarting per chunk, and the
  staging head pools at least four modality slots as released.
- forward accepts temporal_mask (batch, n_patches): padded patches reach
  the head as zero embeddings masked out of its Transformer, as the
  release collates them, and are masked in the encoder too.
- The tokenizer takes a padding mask. In training it only computes the
  valid patches, so masked channels and padded patches affect neither the
  BatchNorm batch statistics nor its running averages (the official
  pretraining normalises its zero padding with the signal; both agree
  when nothing is masked and always in eval mode). This also fixes
  SleepFM.encode.
- from_pretrained completes the tokenizer + head braindecode/SleepFMStager
  mirror with the encoder of the braindecode/SleepFM mirror, keys
  unchanged; checkpoints saved from a SleepFMStager load as they are.
- Shared encoder steps live in _SleepFMSequenceMixin and
  _temporal_transformer, used by SleepFM, the stager and the head.

With released weights, official two-stage pipeline vs public forward:
max |logit diff| 0.0 (float32), input gradients 0.0, BatchNorm running
stats 0.0, in eval and train mode, with masked channels and temporal
padding (train mode compared with a leak-free composition of the
official modules).

* FIX: SleepFM keeps the unmasked training path and tests the pipeline

Review follow-up to the SleepFMStager pipeline fix:
- Without channel_mask or temporal_mask the tokenizer takes its plain
  path again; the masked gather added host syncs and graph breaks to every
  training step (dynamo graph breaks back to 0 for SleepFM, 2 for the
  stager, as before the fix).
- In eval mode masked patches are zeroed before the tokenizer, so a
  non-finite padded channel no longer turns parameter gradients to NaN.
- from_pretrained warns when the released stager is loaded without
  channel_modalities, and accepts encoder_revision for the encoder repo.
- New tests compare the stager with the two-stage pipeline built from
  SleepFM.encode per modality and chunk plus the staging head, with the
  temporal mask aligned to a chunk and ending inside one; they fail when
  either the head's or the encoder's temporal mask is removed.
- Docstrings and changelog state precisely that only masked channels
  deviate from the official BatchNorm behaviour.

* FIX: SleepFMStager warns without channel_modalities for complete checkpoints

from_pretrained only warned about a missing channel_modalities when the
encoder weights were completed from braindecode/SleepFM, so a
self-contained stager checkpoint loaded silently as a single modality.
Warn whenever the loaded model has no channel_modalities, once per call,
and describe the head-only mirror layout as that of older revisions.

* DOC: keep the SleepFM NOTICE entry inside the CC BY-NC file list

* FIX register SleepFM and SleepFMStager under cc-by-nc-4.0

* STY ruff-format sleepfm.py class declarations

* MAINT SleepFM: model classes first, shared steps as functions, PatchTokenizer patching, paper details in docstrings

* MAINT SleepFM: cut redundant checks, mask helpers, loaders and tests; outputs unchanged

* SleepFM: docstring figure, pretraining note and weights box; fold one-use helpers

* SleepFM: note the SHHS-only replication in the docstrings; warn at the caller's line

---------

Co-authored-by: Bruno Aristimunha <[email protected]>
Co-authored-by: Bru <[email protected]>
Co-authored-by: Adam Mounir <[email protected]>