Self-contained stager export (braindecode PR #1106 @ b0ae1aa8)

model.safetensors now holds the complete stager (180 tensors, bit-identical to upstream model_base/best.pt + model_sleep_staging/best.pth); config.json lists every constructor argument (channel_modalities null by design); card updated with a Revisions note. Licence fields unchanged. Previous revision 8681fbade497e2ce103e1b68b01fbb3d56e809f2 still loads.

README.md

@@ -23,15 +23,17 @@ SleepFM is a multimodal polysomnography (PSG) foundation model introduced in:
 
 | File | Description |
 |------|-------------|
-| `model.safetensors` | The full stager, with the parameter names of `braindecode.models.SleepFMStager` |
+| `model.safetensors` | The complete stager (180 tensors: tokenizer, channel pooling, temporal Transformer and staging head), with the parameter names of `braindecode.models.SleepFMStager` |
 | `config.json` | Architecture of the checkpoint, read by `from_pretrained()` |
 
-Upstream ships the stager in two pieces: the channel-agnostic tokenizer lives in the
-encoder checkpoint (`model_base/best.pt`) and the staging head in
-`model_sleep_staging/best.pth`. This file **merges both**, so a single call returns a
-model that is pretrained end to end, its five-class output layer included. The tensors are
-those of the upstream artifacts; only the keys were rewritten to the library's parameter
-names. Loading this file or the two upstream ones gives bit-identical outputs.
+Upstream ships the stager in two pieces: the encoder (channel-agnostic tokenizer,
+channel pooling and temporal Transformer) lives in `model_base/best.pt` and the staging
+head in `model_sleep_staging/best.pth`. This file **merges both**, so a single call
+returns a model that is pretrained end to end, its five-class output layer included. It
+holds the whole encoder except its trial-level temporal pooling, which sleep staging does
+not use. The tensors are those of the upstream artifacts; only the keys were rewritten to
+the library's parameter names. Loading this file or the two upstream ones gives
+bit-identical outputs.
 
 The upstream artifacts themselves are kept, byte-for-byte, in
 [`braindecode/SleepFM`](https://huggingface.co/braindecode/SleepFM).
@@ -41,17 +43,41 @@ The upstream artifacts themselves are kept, byte-for-byte, in
 ```python
 from braindecode.models import SleepFMStager
 
-# Defaults to this repository.
-model = SleepFMStager.from_pretrained(n_chans=4, n_outputs=5, n_times=3840, sfreq=128)
+# Defaults to this repository. The release encodes BAS, RESP, EKG and EMG
+# channels as separate modalities; name the modality of each channel.
+model = SleepFMStager.from_pretrained(
+    n_chans=7,
+    n_outputs=5,
+    n_times=38400,
+    sfreq=128,
+    channel_modalities=["BAS"] * 3 + ["RESP"] * 2 + ["EKG", "EMG"],
+)
 model.eval()
 ```
 
+`config.json` leaves `channel_modalities` unset (`null`), because it depends on the
+montage. Without it every channel is encoded as a single modality and `from_pretrained`
+warns.
+
 The output has shape `(batch, n_outputs, n_patches)`: one prediction per **5-second
 patch**, not per 30-second scoring epoch, so six predictions cover one scored epoch. For
 this checkpoint the five classes are Wake, N1, N2, N3 and REM. Input must be sampled at
 **128 Hz**. Pass `n_outputs` different from 5 to reinitialise the output layer for another
 label set.
 
+## Revisions
+
+- [`8681fba`](https://huggingface.co/braindecode/SleepFMStager/tree/8681fbade497e2ce103e1b68b01fbb3d56e809f2):
+  tokenizer and staging head only (93 tensors). `SleepFMStager.from_pretrained` reads the
+  encoder's channel pooling and temporal Transformer from
+  [`braindecode/SleepFM`](https://huggingface.co/braindecode/SleepFM) at load time.
+- Current revision: the complete stager (180 tensors), bit-identical to the upstream
+  `model_base/best.pt` + `model_sleep_staging/best.pth`, re-exported with the fixed port of
+  braindecode PR #1106. `config.json` now lists every constructor argument; the defaults
+  (`n_chans=4`, `n_times=3840`, `sfreq=128`) are unchanged. Its outputs are bit-identical to
+  those of `8681fba` completed from `braindecode/SleepFM`, and both revisions load with
+  `SleepFMStager.from_pretrained` (pass `revision=` to pin one).
+
 ## License & attribution
 
 - **License: Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0).**

config.json

@@ -5,12 +5,19 @@
   "n_times": 3840,
   "input_window_seconds": null,
   "sfreq": 128.0,
+  "channel_modalities": null,
   "patch_size": 640,
   "embed_dim": 128,
+  "encoder_num_heads": 8,
+  "encoder_num_layers": 6,
+  "encoder_pooling_heads": 8,
+  "encoder_drop_prob": 0.0,
+  "encoder_chunk_patches": 60,
   "staging_num_heads": 4,
   "staging_num_layers": 1,
   "staging_pooling_heads": 4,
   "drop_prob": 0.3,
   "max_seq_length": 8196,
-  "activation": "torch.nn.modules.activation.ELU"
+  "activation": "torch.nn.modules.activation.ELU",
+  "braindecode_version": "1.8.1"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b5caa06bad848f6b2b1c200263ba098954c920cc4505b411182ec47cb5888cd9
-size 9702172
+oid sha256:6a0cdbee2fc6eaf5383cc85984f0b30f4911ac8ca6254ddd5bb1fbecdad3039a
+size 26383332