SleepJEPA

Learning the latent world of sleep!

Overview

SleepJEPA is a Python package for at-home sleep study data to classify sleep stages, sleepiness, and estimate long-term disease risk. SleepJEPA was trained with the JEPA architecture, modified for signal data and sleep studies. The repository is built using nbdev, which means the package is developed in Jupyter notebooks.

See the publication in tbd…

Install

pip install sleepjepa

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 sleepjepa package should be made by editing these notebooks.

To build the package, run nbdev_prepare in the terminal. This will generate the sleepjepa package in the sleepjepa/ 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 sleepjepa 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
  • human_sleep_project/: Processing scripts for the human sleep project dataset
  • jepa/: Training script for the JEPA sleep model
  • mesa/: Processing scripts for the mesa dataset
  • mnc/: Processing scripts for the mnc dataset
  • mros/: Processing scripts for the mros dataset
  • shhs/: Processing scripts for the shhs dataset
  • wsc/: Processing scripts for the wsc dataset
  • sleep_outcomes/:
    • config/: Config files for training outcomes/models
    • train_age.py: Model training script for SleepJEPA representations to estimate chronological age prediction
    • train_obj_sleepiness.py: Model training script for SleepJEPA representations to estimate objective sleepiness and narcolepsy outcomes
    • train_sleep_stages.py: Model training script for SleepJEPA representations to classify sleep stages
    • train_demographics_long_term_outcome.py: Model training script to train demographics only model for long term disease risk estimation
    • train_sleep_long_term_outcome.py: Model training script to train SleepJEPA representations +/- demographicss for long term disease risk estimation

Technical Details

• We trained the foundational model on H100nvl GPUs using PyTorch Lightning.
• We monitored training using the Weights and Biases platform.