PFTSleep

Overview

PFTSleep is a Python package for sleep stage classification using a pre-trained foundational transformer. The repository is built using nbdev, which means the package is developed in Jupyter notebooks.

See the publication in SLEEP and original preprint for more details.

Install

pip install PFTSleep

Repository Structure and Usage

This is an nbdev repository, which means the package is developed in Jupyter notebooks located in the nbs/ directory. Any modifications or additions to the PFTSleep package should be made by editing these notebooks.

To build the package, run nbdev_prepare in the terminal. This will generate the PFTSleep package in the PFTSleep/ directory and all python modules, which can be imported and used in other Python projects.

To add new functionality, create a new notebook or add to exisitng in the nbs/ directory and follow the instructions in the nbdev documentation to add the new functionality. Then, run nbdev_prepare to generate the PFTSleep package with the new functionality.

Directory Structure:

  • nbs/: Contains the source notebooks that generate the Python package
  • jobs/: Contains processing and training scripts
    • apples/: Processing scripts for the apples dataset
    • mesa/: Processing scripts for the mesa dataset
    • mros/: Processing scripts for the mros dataset
    • shhs/: Processing scripts for the shhs dataset
    • wsc/: Processing scripts for the wsc dataset
    • model_training/:
      • train_transformer.py: Trains the initial foundational transformer model
      • train_classifier.py: Trains the probing head for sleep stage classification

Each dataset directory contains scripts to: - Create hypnogram CSVs from annotations - Build zarr files from EDF files - Process and standardize the data for model training

Model Training Pipeline

  1. Foundation Model Training (jobs/model_training/train_transformer.py)
    • Trains the base transformer model on sleep data zarr files
    • Creates general purpose representations of sleep signals
  2. Probe Training (jobs/model_training/train_classifier.py)
    • Trains a classification head on top of the foundation model

Technical Details

  • We trained the foundational model on 2x H100 80gb GPUs using PyTorch Lightning.
  • We monitored training using the Weights and Biases platform.
  • We performed hyperparameter optimization using Optuna.

Citation

If you use PFTSleep in your research, please cite:

@ARTICLE{Fox2025-zc,
  title     = "A foundational transformer leveraging full night, multichannel
               sleep study data accurately classifies sleep stages",
  author    = "Fox, Benjamin and Jiang, Joy and Wickramaratne, Sajila and
               Kovatch, Patricia and Suarez-Farinas, Mayte and Shah, Neomi A
               and Parekh, Ankit and Nadkarni, Girish N",
  journal   = "Sleep",
  publisher = "Oxford University Press (OUP)",
  month     =  mar,
  year      =  2025,
  language  = "en"
}