Improve model card metadata and add research links
Browse filesHi there,
I'm Niels from the Hugging Face community science team. I've updated your model card to improve its visibility and provide more context for researchers:
- Updated the **metadata** to use the `other` pipeline tag as requested for foundation models and ensured the `library_name` is correctly identified.
- Added direct links to the **paper** on Hugging Face and the **GitHub repository**.
- Retained the **Quick Start** code snippet and **Input Format** documentation found in your repository to help users get started quickly.
These changes will help the community better discover and utilize the OSF foundation models!
README.md
CHANGED
|
@@ -1,27 +1,29 @@
|
|
| 1 |
---
|
|
|
|
|
|
|
|
|
|
| 2 |
license: mit
|
|
|
|
| 3 |
tags:
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
pipeline_tag: feature-extraction
|
| 12 |
-
library_name: pytorch
|
| 13 |
-
language:
|
| 14 |
-
- en
|
| 15 |
---
|
| 16 |
|
| 17 |
# OSF: On Pre-training and Scaling of Sleep Foundation Models
|
| 18 |
|
| 19 |
-
[
|
|
|
|
|
|
|
| 20 |
[](https://yang-ai-lab.github.io/osf/)
|
|
|
|
| 21 |
[](LICENSE)
|
| 22 |
[](#installation)
|
| 23 |
|
| 24 |
-
|
| 25 |
## π₯ News
|
| 26 |
|
| 27 |
- [2026-2-24] Our codebase and checkpoint is released. Full codebase for benchmarking will be public available after acceptance.
|
|
@@ -29,26 +31,7 @@ language:
|
|
| 29 |
|
| 30 |
## π Introduction
|
| 31 |
|
| 32 |
-
Polysomnography (PSG) provides the gold standard for sleep assessment but suffers from substantial heterogeneity across recording devices and cohorts.
|
| 33 |
-
There have been growing efforts to build general-purpose foundation models (FMs) for sleep physiology, but lack an in-depth understanding of the pre-training process and scaling patterns that lead to more generalizable sleep FMs.
|
| 34 |
-
To fill this gap, we curate a massive corpus of 166,500 hours of sleep recordings from nine public sources and establish SleepBench, a comprehensive, fully open-source benchmark.
|
| 35 |
-
Leveraging SleepBench, we systematically evaluate four families of self-supervised pre-training objectives and uncover three critical findings:
|
| 36 |
-
(1) existing FMs fail to generalize to missing channels at inference;
|
| 37 |
-
(2) channel-invariant feature learning is essential for pre-training;
|
| 38 |
-
and (3) scaling sample size, model capacity, and multi-source data mixture consistently improves downstream performance.
|
| 39 |
-
With an enhanced pre-training and scaling recipe, we introduce OSF, a family of sleep FMs that achieves state-of-the-art performance across nine datasets on diverse sleep and disease prediction tasks.
|
| 40 |
-
Further analysis of OSF also reveals intriguing properties in sample efficiency, hierarchical aggregation, and cross-dataset scaling.
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
## π Table of Contents
|
| 44 |
-
|
| 45 |
-
1. [Installation](#-installation)
|
| 46 |
-
2. [Quick Start](#-quick-start)
|
| 47 |
-
3. [Pretrained Weights](#-pretrained-weights)
|
| 48 |
-
4. [Usage](#-usage)
|
| 49 |
-
5. [Benchmark Evaluations](#-benchmark-evaluations)
|
| 50 |
-
6. [Supported Datasets](#-supported-datasets)
|
| 51 |
-
7. [Citation](#-citation)
|
| 52 |
|
| 53 |
## πΏ Installation
|
| 54 |
|
|
@@ -59,14 +42,12 @@ conda env create -f environment.yml
|
|
| 59 |
conda activate myenv
|
| 60 |
```
|
| 61 |
|
| 62 |
-
|
| 63 |
### Dependencies
|
| 64 |
|
| 65 |
- Python >= 3.10
|
| 66 |
- PyTorch >= 2.9.0
|
| 67 |
- PyTorch Lightning >= 2.5.5
|
| 68 |
|
| 69 |
-
|
| 70 |
## π Quick Start
|
| 71 |
|
| 72 |
We provide a demo notebook (`demo.ipynb`) demonstrating how to extract embeddings from PSG signals using the pretrained model.
|
|
@@ -98,25 +79,6 @@ with torch.no_grad():
|
|
| 98 |
# patch_embs: [B, 90, 768] - Local patch representations
|
| 99 |
```
|
| 100 |
|
| 101 |
-
## π¦ Pretrained Weights
|
| 102 |
-
|
| 103 |
-
| Model | Backbone | Channels |
|
| 104 |
-
|-------|----------|----------|
|
| 105 |
-
| OSF | ViT-Base | 12-ch |
|
| 106 |
-
|
| 107 |
-
The pretrained weights are included in this repository. You can download them via the Hugging Face Hub:
|
| 108 |
-
|
| 109 |
-
```python
|
| 110 |
-
from huggingface_hub import hf_hub_download
|
| 111 |
-
checkpoint_path = hf_hub_download(repo_id="yang-ai-lab/OSF-Base", filename="osf_backbone.pth")
|
| 112 |
-
```
|
| 113 |
-
|
| 114 |
-
Or via the CLI:
|
| 115 |
-
|
| 116 |
-
```bash
|
| 117 |
-
huggingface-cli download yang-ai-lab/OSF-Base osf_backbone.pth
|
| 118 |
-
```
|
| 119 |
-
|
| 120 |
## π©βπ» Usage
|
| 121 |
|
| 122 |
### Input Format
|
|
@@ -127,130 +89,23 @@ Expected input format:
|
|
| 127 |
- **Epoch Length**: 30 seconds
|
| 128 |
- **Input Shape**: `[B, 12, 1920]`
|
| 129 |
|
| 130 |
-
### Pretraining
|
| 131 |
-
|
| 132 |
-
We support multiple self-supervised pretraining methods, for example, to launch pre-training of our OSF method, run pretraining:
|
| 133 |
-
|
| 134 |
-
```bash
|
| 135 |
-
python main_pretrain.py \
|
| 136 |
-
--model_name "dino_ours" \
|
| 137 |
-
--psg_encoder_name "vit_base" \
|
| 138 |
-
--batch_size 256 \
|
| 139 |
-
--lr 5e-5 \
|
| 140 |
-
--max_epochs 30 \
|
| 141 |
-
--num_devices 4 \
|
| 142 |
-
--patch_size_time 64 \
|
| 143 |
-
--patch_size_ch 4 \
|
| 144 |
-
--precision "bf16-mixed"
|
| 145 |
-
```
|
| 146 |
-
|
| 147 |
-
See `main_pipleines/main_pretrain.py` for more detailed settings.
|
| 148 |
-
|
| 149 |
-
### Fine-tuning
|
| 150 |
-
|
| 151 |
-
Fine-tune the pretrained model on downstream tasks:
|
| 152 |
-
|
| 153 |
-
```bash
|
| 154 |
-
python main_finetune.py \
|
| 155 |
-
--model_name "dino_ours" \
|
| 156 |
-
--ckpt_path "/path/to/pretrained/checkpoint.ckpt" \
|
| 157 |
-
--downstream_dataset_name "shhs" \
|
| 158 |
-
--eval_label "Stage" \
|
| 159 |
-
--train_data_pct 1.0 \
|
| 160 |
-
--max_steps 500 \
|
| 161 |
-
--lr 0.1 \
|
| 162 |
-
--num_devices 4
|
| 163 |
-
```
|
| 164 |
|
|
|
|
| 165 |
|
| 166 |
## π Benchmark Evaluations
|
| 167 |
|
| 168 |
-
|
| 169 |
-
|
| 170 |
-
| Method | Type | Original Paper |
|
| 171 |
-
|--------|------|-------------|
|
| 172 |
-
| SleepFM | Contrastive | Leave-one-out multi-modal contrastive learning |
|
| 173 |
-
| SimCLR | Contrastive | Simple Contrastive Learning |
|
| 174 |
-
| DINO | Self-distillation | DINOv2 (Oquab et al., 2023) |
|
| 175 |
-
| VQ-VAE | Reconstruction | Vector-quantized variational autoencoder |
|
| 176 |
-
| MAE | Reconstruction | Masked Autoencoding |
|
| 177 |
-
| AR | Autoregressive | Autoregressive Next-Token prediction |
|
| 178 |
-
| OSF | Self-distillation | ours |
|
| 179 |
-
|
| 180 |
-
### Downstream Tasks
|
| 181 |
-
|
| 182 |
-
**Epoch-level Classification Tasks:**
|
| 183 |
-
|
| 184 |
-
| Task | Classes | Description |
|
| 185 |
-
|------|---------|-------------|
|
| 186 |
-
| Sleep Stage | 4 | Awake, Light Sleep, Deep Sleep, REM classification |
|
| 187 |
-
| Arousal | 2 | Arousal event detection |
|
| 188 |
-
| Hypopnea | 2 | Hypopnea event detection |
|
| 189 |
-
| Oxygen Desaturation | 2 | Oxygen desaturation detection |
|
| 190 |
-
|
| 191 |
-
|
| 192 |
-
### Evaluation Settings
|
| 193 |
-
|
| 194 |
-
| Setting | Description |
|
| 195 |
-
|---------|-------------|
|
| 196 |
-
| Linear Probing | Freeze backbone, train linear classifier |
|
| 197 |
-
| Full Fine-tuning | Fine-tune entire model end-to-end |
|
| 198 |
-
| Few-shot (k-shot) | Train with limited labeled samples |
|
| 199 |
-
|
| 200 |
-
For example scripts, see `main_pipelines` and `bash_scripts` folders.
|
| 201 |
-
|
| 202 |
-
## π Supported Datasets
|
| 203 |
-
|
| 204 |
-
We aggregated nine large-scale datasets from the National Sleep Research Resource platform.
|
| 205 |
-
|
| 206 |
-
| Dataset | Full Name | Source |
|
| 207 |
-
|---------|-----------|--------|
|
| 208 |
-
| SHHS | Sleep Heart Health Study | NSRR |
|
| 209 |
-
| CHAT | Childhood Adenotonsillectomy Trial | NSRR |
|
| 210 |
-
| MROS | MrOS Sleep Study | NSRR |
|
| 211 |
-
| CCSHS | Cleveland Children's Sleep and Health Study | NSRR |
|
| 212 |
-
| CFS | Cleveland Family Study | NSRR |
|
| 213 |
-
| MESA | Multi-Ethnic Study of Atherosclerosis | NSRR |
|
| 214 |
-
| SOF | Study of Osteoporotic Fractures | NSRR |
|
| 215 |
-
| WSC | Wisconsin Sleep Cohort | NSRR |
|
| 216 |
-
| STAGES | Stanford Technology Analytics and Genomics in Sleep | NSRR |
|
| 217 |
-
| NCHSDB | NCH Sleep DataBank | NSRR |
|
| 218 |
-
|
| 219 |
-
For new users, please apply for an account and access to each of these datasets following instructions here [NSRR Registration](https://sleepdata.org/join)
|
| 220 |
-
|
| 221 |
-
## π Project Structure
|
| 222 |
-
|
| 223 |
-
```
|
| 224 |
-
OSF-Open-Sleep-Foundation-Model/
|
| 225 |
-
βββ osf/
|
| 226 |
-
β βββ backbone/ # ViT backbone implementations
|
| 227 |
-
β β βββ vit1d_cls.py
|
| 228 |
-
β βββ models/ # SSL model implementations
|
| 229 |
-
β β βββ dino_model_cls.py
|
| 230 |
-
β β
|
| 231 |
-
β βββ datasets/ # Data loading utilities
|
| 232 |
-
β βββ utils/ # Helper functions
|
| 233 |
-
βββ main_pipelines/ # Training scripts
|
| 234 |
-
β βββ main_pretrain.py
|
| 235 |
-
β βββ ...
|
| 236 |
-
βββ bash_scripts/ # Example bash scripts
|
| 237 |
-
βββ osf_backbone.pth # Pretrained model weights
|
| 238 |
-
βββ demo.ipynb # Quick start demo
|
| 239 |
-
βββ config.py # Dataset and channel configurations
|
| 240 |
-
βββ train_config.py # Training configurations
|
| 241 |
-
```
|
| 242 |
-
|
| 243 |
|
| 244 |
## π Citation
|
| 245 |
|
| 246 |
-
If you use this code or models in your research, please cite
|
| 247 |
|
| 248 |
```bibtex
|
| 249 |
@article{shuai2026osf,
|
| 250 |
title={OSF: On Pre-training and Scaling of Sleep Foundation Models},
|
| 251 |
author={Shuai, Zitao and Xu, Zongzhe and Yang, David and Wang, Wei and Yang, Yuzhe},
|
| 252 |
-
journal={arXiv preprint},
|
| 253 |
year={2026}
|
| 254 |
}
|
| 255 |
-
```
|
| 256 |
-
|
|
|
|
| 1 |
---
|
| 2 |
+
language:
|
| 3 |
+
- en
|
| 4 |
+
library_name: pytorch
|
| 5 |
license: mit
|
| 6 |
+
pipeline_tag: other
|
| 7 |
tags:
|
| 8 |
+
- sleep
|
| 9 |
+
- eeg
|
| 10 |
+
- polysomnography
|
| 11 |
+
- foundation-model
|
| 12 |
+
- self-supervised
|
| 13 |
+
- vit
|
| 14 |
+
- biosignals
|
|
|
|
|
|
|
|
|
|
|
|
|
| 15 |
---
|
| 16 |
|
| 17 |
# OSF: On Pre-training and Scaling of Sleep Foundation Models
|
| 18 |
|
| 19 |
+
This repository contains the weights for **OSF**, a family of sleep foundation models introduced in the paper [OSF: On Pre-training and Scaling of Sleep Foundation Models](https://huggingface.co/papers/2603.00190).
|
| 20 |
+
|
| 21 |
+
[](https://huggingface.co/papers/2603.00190)
|
| 22 |
[](https://yang-ai-lab.github.io/osf/)
|
| 23 |
+
[](https://github.com/yang-ai-lab/OSF-Open-Sleep-FM)
|
| 24 |
[](LICENSE)
|
| 25 |
[](#installation)
|
| 26 |
|
|
|
|
| 27 |
## π₯ News
|
| 28 |
|
| 29 |
- [2026-2-24] Our codebase and checkpoint is released. Full codebase for benchmarking will be public available after acceptance.
|
|
|
|
| 31 |
|
| 32 |
## π Introduction
|
| 33 |
|
| 34 |
+
Polysomnography (PSG) provides the gold standard for sleep assessment but suffers from substantial heterogeneity across recording devices and cohorts. OSF is a family of sleep foundation models (FMs) pre-trained on a massive corpus of 166,500 hours of sleep recordings from nine public sources. Leveraging the SleepBench benchmark, the authors establish an enhanced pre-training and scaling recipe that achieves state-of-the-art performance across diverse sleep and disease prediction tasks.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 35 |
|
| 36 |
## πΏ Installation
|
| 37 |
|
|
|
|
| 42 |
conda activate myenv
|
| 43 |
```
|
| 44 |
|
|
|
|
| 45 |
### Dependencies
|
| 46 |
|
| 47 |
- Python >= 3.10
|
| 48 |
- PyTorch >= 2.9.0
|
| 49 |
- PyTorch Lightning >= 2.5.5
|
| 50 |
|
|
|
|
| 51 |
## π Quick Start
|
| 52 |
|
| 53 |
We provide a demo notebook (`demo.ipynb`) demonstrating how to extract embeddings from PSG signals using the pretrained model.
|
|
|
|
| 79 |
# patch_embs: [B, 90, 768] - Local patch representations
|
| 80 |
```
|
| 81 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 82 |
## π©βπ» Usage
|
| 83 |
|
| 84 |
### Input Format
|
|
|
|
| 89 |
- **Epoch Length**: 30 seconds
|
| 90 |
- **Input Shape**: `[B, 12, 1920]`
|
| 91 |
|
| 92 |
+
### Pretraining and Fine-tuning
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 93 |
|
| 94 |
+
For detailed instructions on pretraining and fine-tuning using the OSF framework, please refer to the scripts in the [official GitHub repository](https://github.com/yang-ai-lab/OSF-Open-Sleep-FM).
|
| 95 |
|
| 96 |
## π Benchmark Evaluations
|
| 97 |
|
| 98 |
+
OSF has been evaluated on the **SleepBench** benchmark across tasks such as Sleep Stage classification, Arousal detection, Hypopnea event detection, and Oxygen Desaturation detection, outperforming existing SSL methods like SleepFM, SimCLR, and DINO.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 99 |
|
| 100 |
## π Citation
|
| 101 |
|
| 102 |
+
If you use this code or models in your research, please cite the paper:
|
| 103 |
|
| 104 |
```bibtex
|
| 105 |
@article{shuai2026osf,
|
| 106 |
title={OSF: On Pre-training and Scaling of Sleep Foundation Models},
|
| 107 |
author={Shuai, Zitao and Xu, Zongzhe and Yang, David and Wang, Wei and Yang, Yuzhe},
|
| 108 |
+
journal={arXiv preprint arXiv:2603.00190},
|
| 109 |
year={2026}
|
| 110 |
}
|
| 111 |
+
```
|
|
|