SleepFM-Clinical Documentation

The disease prediction model (based on a COXPH framework) outputs log-probabilities for various medical conditions.

Output Format: The model typically outputs a vector of size 1065, representing log-probabilities for 1065 different conditions.

Mapping Outputs to Diseases: To understand what each index in the output vector represents, map the indices to their corresponding phecodes using the sleepfm/configs/label_mapping.csv file.

Data Requirements: The DiagnosisFinetuneFullCOXPHWithDemoDataset requires:

• HDF5 files containing modality embeddings.
• Demographic data (e.g., age, gender) in a CSV.
• Event indicators (is_event.csv) and time-to-event data (time_to_event.csv).

# Mapping model outputs to disease labels
import pandas as pd

labels_df = pd.read_csv("../sleepfm/configs/label_mapping.csv")

# Assuming 'all_outputs' are the model's logprobs and 'all_is_event' are the event indicators
labels_df["output"] = all_outputs[0]
labels_df["is_event"] = all_is_event[0]
labels_df["event_time"] = all_event_times[0]

Before using SleepFM, you must convert your PSG (Polysomnography) recordings from EDF format to HDF5. You can do this using the EDFToHDF5Converter class.

Important: PSG recordings often have varying channel sets. You must ensure your channels are correctly mapped to modalities (e.g., EEG, EOG, EMG, ECG) by updating sleepfm/configs/channel_groups.json. If your data contains channels not covered by the default mapping, consult a domain expert to categorize them correctly in the config file.

```python
from preprocessing.preprocessing import EDFToHDF5Converter
import os

base_save_path = "demo_data"
root_dir = "/edf_root"
target_dir = "/note"

edf_path = "demo_data/demo_psg.edf"
hdf5_path = os.path.join(base_save_path, "demo_psg.hdf5")

converter = EDFToHDF5Converter(
    root_dir=root_dir,
    target_dir=target_dir,
    resample_rate=128
)

# Convert a single file
converter.convert(edf_path, hdf5_path)

For batch processing multiple files, use the shell script: sleepfm/preprocessing/preprocessing.sh.

Sleep staging classifies sleep stages (e.g., Wake, Stage 1, Stage 2, Stage 3, REM) from PSG data. While pretrained models are available, it is highly recommended to finetune the model head on your specific dataset to account for distribution shifts.

Finetuning: Use the script sleepfm/pipeline/finetune_sleep_staging.py to adapt the model to your data.

Data Loading: Use SleepEventClassificationDataset which requires HDF5 files and corresponding CSV label files. The dataset supports a context parameter to define the sequence length for temporal modeling.

from models.dataset import SleepEventClassificationDataset
from torch.utils.data import DataLoader

# Initialize dataset
test_dataset = SleepEventClassificationDataset(
    sleep_staging_config, 
    channel_groups, 
    split="test", 
    hdf5_paths=hdf5_paths, 
    label_files=label_files
)

# Use the specific collate function for sleep staging
test_loader = DataLoader(
    test_dataset, 
    batch_size=8, 
    shuffle=False, 
    num_workers=1, 
    collate_fn=sleep_event_finetune_full_collate_fn
)

You can generate two types of embeddings from the SleepFM pretrained model:

1. Granular 5-second-level embeddings: High-resolution temporal embeddings.
2. Aggregated 5-minute-level embeddings: Lower-resolution embeddings suitable for long-term analysis.

To perform this manually, load the model using its config.json and best.pt checkpoint, then pass the data through a SetTransformerDataset and a DataLoader using the collate_fn.

# For full automation, use the pipeline script:
# sleepfm/pipeline/generate_embeddings.py

# Manual snippet for model initialization:
import torch
import sys
from utils import load_config, load_data

model_path = "../sleepfm/checkpoints/model_base"
config_path = os.path.join(model_path, "config.json")
config = load_config(config_path)

# Dynamically instantiate the model class defined in config
model_class = getattr(sys.modules[__name__], config['model'])
model = model_class(
    config['in_channels'], 
    config['patch_size'], 
    config['embed_dim'], 
    config['num_heads'], 
    config['num_layers'], 
    pooling_head=config['pooling_head'], 
    dropout=0.0
)

checkpoint = torch.load(os.path.join(model_path, "best.pt"))
model.load_state_dict(checkpoint["state_dict"])
model.eval()