Add the braindecode-format encoder (safetensors + config)

README.md

@@ -9,9 +9,9 @@ tags:
   - braindecode
 ---
 
-# SleepFM — pretrained checkpoints
+# SleepFM — pretrained encoder
 
-Mirror of the official **SleepFM** checkpoints, re-hosted for stable loading from
+Mirror of the official **SleepFM** encoder checkpoint, re-hosted for stable loading from
 [Braindecode](https://github.com/braindecode/braindecode).
 
 SleepFM is a multimodal polysomnography (PSG) foundation model introduced in:
@@ -19,37 +19,40 @@ SleepFM is a multimodal polysomnography (PSG) foundation model introduced in:
 > R. Thapa et al., *"A multimodal sleep foundation model for disease prediction,"*
 > **Nature Medicine** (2026). https://doi.org/10.1038/s41591-025-04133-4
 
+The downstream sleep stager lives in a separate repository,
+[`braindecode/SleepFMStager`](https://huggingface.co/braindecode/SleepFMStager),
+because a Braindecode `config.json` describes exactly one architecture.
+
 ## Files
 
-| File | Description | Used by |
-|------|-------------|---------|
-| `model_base/best.pt` | Pretrained channel-agnostic PSG encoder (backbone) | `SleepFM.load_pretrained_backbone`, `SleepFMStager.load_pretrained_backbone` |
-| `model_sleep_staging/best.pth` | Downstream sleep-staging head | `SleepFMStager.load_pretrained_staging_head` |
+| File | Description |
+|------|-------------|
+| `model.safetensors` | The encoder, with the parameter names of `braindecode.models.SleepFM` |
+| `config.json` | Architecture of the checkpoint, read by `from_pretrained()` |
+| `model_base/best.pt` | The upstream artifact, byte-for-byte, kept for provenance |
+| `model_sleep_staging/best.pth` | The upstream staging artifact, byte-for-byte (see `SleepFMStager`) |
+
+`model.safetensors` holds the **same tensors** as `model_base/best.pt`; only the keys were
+rewritten (the `module.` prefix of the distributed training run stripped, and
+`positional_encoding.pe` renamed) so that the library needs no remapping code at load
+time. Loading either way gives bit-identical outputs.
 
-These are byte-for-byte copies of the upstream artifacts; only the hosting location
-changed. The layout (`model_base/`, `model_sleep_staging/`) mirrors the upstream repo.
+Note that the released encoder is contrastive and carries **no classification head**:
+`final_layer` is randomly initialised and must be fine-tuned.
 
 ## Usage
 
 ```python
-import torch
-from braindecode.models import SleepFMStager
-
-base = torch.hub.load_state_dict_from_url(
-    "https://huggingface.co/braindecode/SleepFM/resolve/main/model_base/best.pt",
-    map_location="cpu",
-)
-staging = torch.hub.load_state_dict_from_url(
-    "https://huggingface.co/braindecode/SleepFM/resolve/main/model_sleep_staging/best.pth",
-    map_location="cpu",
-)
-
-model = SleepFMStager(n_chans=4, n_outputs=5, n_times=3840, sfreq=128)
-model.load_pretrained_backbone(base)
-model.load_pretrained_staging_head(staging)
+from braindecode.models import SleepFM
+
+# Defaults to this repository.
+model = SleepFM.from_pretrained(n_chans=4, n_outputs=5, n_times=3840, sfreq=128)
 model.eval()
 ```
 
+Input must be sampled at **128 Hz**; the reference `patch_size=640` is a 5-second patch at
+that rate. A channel mask of shape `(batch, n_chans)` marks missing channels with `True`.
+
 ## License & attribution
 
 - **License: Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0).**

config.json

@@ -0,0 +1,16 @@
+{
+  "n_outputs": 5,
+  "n_chans": 4,
+  "chs_info": null,
+  "n_times": 3840,
+  "input_window_seconds": null,
+  "sfreq": 128.0,
+  "patch_size": 640,
+  "embed_dim": 128,
+  "num_heads": 8,
+  "num_layers": 6,
+  "pooling_heads": 8,
+  "drop_prob": 0.3,
+  "max_seq_length": 128,
+  "activation": "torch.nn.modules.activation.ELU"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a8d02f16323e564b834f3f31ecf9995f4401ff1598459feaec7f3ab671366b2
+size 19410992