SleepFM-Clinical

https://github.com/zou-group/sleepfm-clinical
Admin
SleepFM-Clinical is a multimodal sleep foundation model that analyzes polysomnography (PSG)...
Tokens:12,221
Snippets:85
Trust Score:9
Update:3 months ago
Show doc for...
Tokens:
Raw
### Clone Repository and Set Up Conda Environment

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/README.md

Clone the SleepFM repository and create a Conda environment with all necessary dependencies. This is the initial setup step for running the project.

```bash
git clone https://github.com/zou-group/sleepfm-clinical.git
cd sleepfm-clinical
conda env create -f env.yml
conda activate sleepfm_env
```

--------------------------------

### Import Core Libraries

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Imports essential Python libraries for data manipulation, machine learning, and deep learning. Ensure these libraries are installed.

```python
import torch
from torch import nn
import numpy as np
from sklearn.metrics import f1_score, confusion_matrix
import os
import tqdm
import random
import matplotlib.pyplot as plt
import seaborn as sns
import sys
from collections import Counter
import pandas as pd
sys.path.append("..")
sys.path.append("../sleepfm")
from preprocessing.preprocessing import EDFToHDF5Converter
from models.dataset import SetTransformerDataset, collate_fn
from models.models import SetTransformer, SleepEventLSTMClassifier, DiagnosisFinetuneFullLSTMCOXPHWithDemo
import h5py
from utils import load_config, load_data, save_data, count_parameters
from torch.utils.data import Dataset, DataLoader
```

--------------------------------

### Output Directory Setup for Embeddings

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Creates output directories for storing generated embeddings and aggregated embeddings. Ensure the base save path is correctly defined.

```python
output = os.path.join(base_save_path, "demo_emb")
output_5min_agg = os.path.join(base_save_path, "demo_5min_agg_emb")
os.makedirs(output, exist_ok=True)
os.makedirs(output_5min_agg, exist_ok=True)
```

--------------------------------

### Initialize and Print Sleep Staging Model Info

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Initializes the sleep staging model and prints its trainable parameters and number of layers. This is useful for verifying model setup.

```python
print(f"Model initialized: {sleep_staging_model_name}")
total_layers, total_params = count_parameters(sleep_staging_model)
print(f'Trainable parameters: {total_params / 1e6:.2f} million')
print(f'Number of layers: {total_layers}')
```

--------------------------------

### Example Sleep Stage Labels Format

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/README.md

Illustrates the expected format for sleep stage labels, which are used for training and evaluation. Ensure your label files match this structure.

```csv
Start,Stop,StageName,StageNumber
0.0,5190.0,Wake,0
0.0,5190.0,Wake,0,
0.0,5190.0,Wake,0
```

--------------------------------

### Get Item from Dataset

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Retrieves data for a given index from the dataset. It reads HDF5 files, extracts relevant features, converts them to PyTorch tensors, and returns them along with labels and metadata.

```python
def __getitem__(self, idx):
    hdf5_path, tte_event = self.index_map[idx]

    event_time = tte_event["event_time"]
    is_event = tte_event["is_event"]
    demo_feats = tte_event["demo_feats"]

    x_data = []
    with h5py.File(hdf5_path, 'r') as hf:
        dset_names = []
        for dset_name in hf.keys():
            if isinstance(hf[dset_name], h5py.Dataset) and dset_name in self.config["modality_types"]:
                dset_names.append(dset_name)
        
        random.shuffle(dset_names)
        for dataset_name in dset_names:
            x_data.append(hf[dataset_name][:])

    if not x_data:
        # Skip this data point if x_data is empty
        return self.__getitem__((idx + 1) % self.total_len)

    # Convert x_data list to a single numpy array
    x_data = np.array(x_data)

    # Convert x_data to tensor
    x_data = torch.tensor(x_data, dtype=torch.float32)

    event_time = torch.tensor(event_time, dtype=torch.float32)
    is_event = torch.tensor(is_event) 

    demo_feats = torch.tensor(demo_feats, dtype=torch.float32)

    return x_data, event_time, is_event, demo_feats, self.max_channels, self.max_seq_len, hdf5_path
```

--------------------------------

### Get Shape of Flattened Arrays

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Retrieves the shapes of the flattened data arrays after concatenation. Useful for verifying data dimensions.

```python
all_logits_flat.shape, all_outputs_flat.shape, all_targets_flat.shape, all_masks_flat.shape
```

--------------------------------

### Initialize Output Directory

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Creates a directory to store demo data. Ensure the path is writable.

```python
base_save_path = "demo_data"
os.makedirs(base_save_path, exist_ok=True)
```

--------------------------------

### Initialize EDF to HDF5 Converter

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Sets up the converter for EDF files. Dummy root and target directories are provided as they are not used for single file conversion in this demo.

```python
base_save_path = "demo_data"
os.makedirs(base_save_path, exist_ok=True)

root_dir = "/edf_root"      # dummy root not used for a single file conversion
target_dir = "/note"    # dummy target not used for a single file conversion

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
)

# run for single file conversion
converter.convert(edf_path, hdf5_path)
```

--------------------------------

### Initialize Diagnosis Dataset

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Sets up the dataset for disease prediction fine-tuning. Specifies configuration, channel groups, split, HDF5 file paths, and demo labels.

```python
hdf5_paths = [os.path.join(base_save_path, "demo_emb/demo_psg.hdf5")]
demo_labels_path = os.path.join(base_save_path, "demo_age_gender.csv")
config["labels_path"] = base_save_path

test_dataset = DiagnosisFinetuneFullCOXPHWithDemoDataset(config, channel_groups, split="test", hdf5_paths=hdf5_paths, demo_labels_path=demo_labels_path)
```

--------------------------------

### Initialize and Configure SleepFM Model

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Instantiates the SleepFM model based on the extracted parameters and moves it to the appropriate device (GPU if available). It also prints the total trainable parameters and layers.

```python
model_class = getattr(sys.modules[__name__], config['model'])
model = model_class(in_channels, patch_size, embed_dim, num_heads, num_layers, pooling_head=pooling_head, dropout=dropout)

device = torch.device("cuda")
if device.type == "cuda":
    model = torch.nn.DataParallel(model)

model.to(device)
total_layers, total_params = count_parameters(model)
print(f'Trainable parameters: {total_params / 1e6:.2f} million')
print(f'Number of layers: {total_layers}')
```

--------------------------------

### Parallelize Model and Print Info

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Wraps the initialized model with `nn.DataParallel` for multi-GPU training and prints model initialization status, trainable parameters, and layer count. Requires a `count_parameters` function.

```python
model = nn.DataParallel(model)
print(f"Model initialized: {model_name}")
total_layers, total_params = count_parameters(model)
print(f'Trainable parameters: {total_params / 1e6:.2f} million')
print(f'Number of layers: {total_layers}')
```

--------------------------------

### Load Finetuned Sleep Staging Model Configuration

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the configuration for a finetuned sleep staging model from a JSON file. Initializes the model class and parameters based on the loaded configuration.

```python
sleep_staging_model_path = "../sleepfm/checkpoints/model_sleep_staging"
sleep_staging_config = load_data(os.path.join(sleep_staging_model_path, "config.json"))

sleep_staging_model_params = sleep_staging_config['model_params']
sleep_staging_model_class = getattr(sys.modules[__name__], sleep_staging_config['model'])

sleep_staging_model = sleep_staging_model_class(**sleep_staging_model_params).to(device)
sleep_staging_model_name = type(sleep_staging_model).__name__
```

--------------------------------

### Load Model Configuration and Channels

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the configuration and channel group data for the pretrained SleepFM model. Ensure the paths to checkpoints and configs are correct.

```python
model_path = "../sleepfm/checkpoints/model_base"
channel_groups_path = "../sleepfm/configs/channel_groups.json"
config_path = os.path.join(model_path, "config.json")

config = load_config(config_path)
channel_groups = load_data(channel_groups_path)
```

--------------------------------

### Configuration Object

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Displays the configuration object used in the notebook. This is useful for understanding the parameters and settings applied to the model or process.

```python
config
```

--------------------------------

### Create Save Directory

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Creates a directory to save prediction outputs. Ensures the directory exists, creating it if necessary.

```python
save_path = os.path.join(base_save_path, "demo_diagnosis")
os.makedirs(save_path, exist_ok=True)
```

--------------------------------

### Define Class Labels and Mapping

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Sets up a list of class labels and creates a dictionary mapping these labels to their corresponding integer indices. This is crucial for consistent interpretation of model outputs and targets.

```python
class_labels = ["Wake", "Stage 1", "Stage 2", "Stage 3", "REM"]
# class_labels = ["No-Apnea", "Apnea"]
class_mapping = {label: idx for idx, label in enumerate(class_labels)}
```

--------------------------------

### Initialize Disease Prediction Model

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Configures and initializes the disease prediction model based on loaded parameters. It dynamically retrieves the model class and moves the model to the appropriate device.

```python
config["model_params"]["dropout"] = 0.0
model_params = config['model_params']
model_class = getattr(sys.modules[__name__], config['model'])
model = model_class(**model_params).to(device)
model_name = type(model).__name__
```

--------------------------------

### Specify Checkpoint Path

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Defines the file path to the pre-trained model checkpoint. Ensure the path is correct.

```python
checkpoint_path = os.path.join(disease_model_path, "best.pth")
```

--------------------------------

### Create DataLoader

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Initializes a DataLoader for the test dataset. Configures batch size, shuffling, number of workers, and a custom collate function.

```python
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=1, collate_fn=diagnosis_finetune_full_coxph_with_demo_collate_fn)
```

--------------------------------

### Load Model State Dictionary

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the saved state dictionary from a checkpoint file into the initialized model. This step is essential for resuming training or performing inference with pre-trained weights.

```python
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint)
```

--------------------------------

### Load Pretrained Model Weights

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the best performing weights from a checkpoint file into the initialized SleepFM model and sets the model to evaluation mode. This prepares the model for inference.

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

--------------------------------

### Model Configuration Parameters

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Configuration parameters for the SetTransformer model, including dataset paths, model dimensions, and training settings.

```python
Result:
{'seed': 42,
 'model': 'SetTransformer',
 'in_channels': 1,
 'batch_size': 128,
 'epochs': 1,
 'lr': 0.001,
 'lr_step_period': 2,
 'gamma': 0.1,
 'temperature': 0.0,
 'momentum': 0.9,
 'num_workers': 16,
 'embed_dim': 128,
 'num_heads': 8,
 'num_layers': 6,
 'pooling_head': 8,
 'dropout': 0.3,
 'split_path': 'path_to_/dataset_split.json',
 'save_path': 'path_to_/models',
 'weight_decay': 0.0,
 'mode': 'leave_one_out',
 'save_iter': 5000,
 'eval_iter': 5000,
 'log_interval': 100,
 'use_wandb': True,
 'BAS_CHANNELS': 10,
 'RESP_CHANNELS': 7,
 'EKG_CHANNELS': 2,
 'EMG_CHANNELS': 4,
 'max_files': None,
 'val_size': 100,
 'sampling_duration': 5,
 'sampling_freq': 128,
 'patch_size': 640,
 'modality_types': ['BAS', 'RESP', 'EKG', 'EMG']}
```

--------------------------------

### Load Sleep Staging Model Checkpoint

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the best performing checkpoint for the sleep staging model. Ensure the model architecture is defined before loading weights.

```python
sleep_staging_checkpoint_path = os.path.join(sleep_staging_model_path, "best.pth")
sleep_staging_checkpoint = torch.load(sleep_staging_checkpoint_path)
sleep_staging_model.load_state_dict(sleep_staging_checkpoint)
```

--------------------------------

### Initialize Sleep Event Classification Dataset

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

This snippet shows how to initialize a `SleepEventClassificationDataset` for testing. It specifies HDF5 paths and label files, and sets the split to 'test'. Ensure `sleep_staging_config` and `channel_groups` are defined prior to use.

```python
hdf5_paths = [os.path.join(base_save_path, "demo_emb/demo_psg.hdf5")]
label_files = [os.path.join(base_save_path, "demo_psg.csv")]
test_dataset = SleepEventClassificationDataset(sleep_staging_config, channel_groups, split="test", hdf5_paths=hdf5_paths, label_files=label_files)
```

--------------------------------

### Create DataLoader for Sleep Staging

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

This code creates a PyTorch DataLoader for the sleep staging dataset. It uses a batch size of 8, disables shuffling, and specifies the custom `sleep_event_finetune_full_collate_fn` for batch collation. Set `num_workers` according to your system's capabilities.

```python
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=1, collate_fn=sleep_event_finetune_full_collate_fn)
```

--------------------------------

### Parallelize Sleep Staging Model

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Wraps the sleep staging model with DataParallel for multi-GPU training and prints the number of GPUs being used.

```python
sleep_staging_model = nn.DataParallel(sleep_staging_model)
print(f"Using {torch.cuda.device_count()} GPUs")
```

--------------------------------

### Display Mapped Labels

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Shows the first few rows of the labels DataFrame, now augmented with model predictions, event status, and event times.

```python
labels_df.head()
```

--------------------------------

### Load SleepFM Dataset and DataLoader

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the dataset using SetTransformerDataset and creates a DataLoader for batch processing. Ensure the base save path and dataset configuration are correctly set.

```python
hdf5_paths = [os.path.join(base_save_path, "demo_psg.hdf5")]
dataset = SetTransformerDataset(config, channel_groups, hdf5_paths=hdf5_paths, split="test")

dataloader = torch.utils.data.DataLoader(dataset, 
                                            batch_size=16, 
                                            num_workers=1, 
                                            shuffle=False, 
                                            collate_fn=collate_fn)
```

--------------------------------

### Load Model Configuration

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the configuration file for the disease prediction model. Assumes the config file is in JSON format.

```python
disease_model_path = "../sleepfm/checkpoints/model_diagnosis"
config = load_data(os.path.join(disease_model_path, "config.json"))
```

--------------------------------

### Extract Model Parameters from Config

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Retrieves model architecture parameters from the loaded configuration. These parameters define the structure of the SleepFM model.

```python
modality_types = config["modality_types"]
in_channels = config["in_channels"]
patch_size = config["patch_size"]
embed_dim = config["embed_dim"]
num_heads = config["num_heads"]
num_layers = config["num_layers"]
pooling_head = config["pooling_head"]
dropout = 0.0
```

--------------------------------

### Collate Function for Diagnosis Finetuning

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Pads and stacks data from a batch of samples for diagnosis finetuning. It handles variable sequence lengths and channel numbers, creating masks for padded regions.

```python
def diagnosis_finetune_full_coxph_with_demo_collate_fn(batch):
    x_data, event_time, is_event, demo_feats, max_channels_list, max_seq_len_list, hdf5_path_list = zip(*batch)

    num_channels = max(max_channels_list)

    if max_seq_len_list[0] == None:
        max_seq_len = max([item.size(1) for item in x_data])
    else:
        max_seq_len = max_seq_len_list[0]

    padded_x_data = []
    padded_mask = []
    for item in x_data:
        c, s, e = item.size()
        c = min(c, num_channels)
        s = min(s, max_seq_len)  # Ensure the sequence length doesn't exceed max_seq_len

        # Create a padded tensor and a mask tensor
        padded_item = torch.zeros((num_channels, max_seq_len, e))
        mask = torch.ones((num_channels, max_seq_len))

        # Copy the actual data to the padded tensor and set the mask for real data
        padded_item[:c, :s, :e] = item[:c, :s, :e]
        mask[:c, :s] = 0  # 0 for real data, 1 for padding

        padded_x_data.append(padded_item)
        padded_mask.append(mask)
    
    # Stack all tensors into a batch
    x_data = torch.stack(padded_x_data)
    event_time = torch.stack(event_time)
    is_event = torch.stack(is_event)
    demo_feats = torch.stack(demo_feats)
    padded_mask = torch.stack(padded_mask)
    
    return x_data, event_time, is_event, demo_feats, padded_mask, hdf5_path_list
```

--------------------------------

### Load Autoreload Extension

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the autoreload extension for interactive development. This is typically used in Jupyter environments.

```python
%load_ext autoreload
%autoreload 2
```

--------------------------------

### SleepFM Model Architecture

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Displays the detailed architecture of the SleepFM model, including its various layers for patch embedding, spatial pooling, and transformer encoding. This is useful for understanding the model's structure and how it processes input data.

```python
Result:
DataParallel(
  (module): SetTransformer(
    (patch_embedding): Tokenizer(
      (tokenizer): Sequential(
        (0): Conv1d(1, 4, kernel_size=(5,), stride=(2,), padding=(2,))
        (1): BatchNorm1d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ELU(alpha=1.0)
        (3): LayerNorm((4, 320), eps=1e-05, elementwise_affine=True)
        (4): Conv1d(4, 8, kernel_size=(5,), stride=(2,), padding=(2,))
        (5): BatchNorm1d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (6): ELU(alpha=1.0)
        (7): LayerNorm((8, 160), eps=1e-05, elementwise_affine=True)
        (8): Conv1d(8, 16, kernel_size=(5,), stride=(2,), padding=(2,))
        (9): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (10): ELU(alpha=1.0)
        (11): LayerNorm((16, 80), eps=1e-05, elementwise_affine=True)
        (12): Conv1d(16, 32, kernel_size=(5,), stride=(2,), padding=(2,))
        (13): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (14): ELU(alpha=1.0)
        (15): LayerNorm((32, 40), eps=1e-05, elementwise_affine=True)
        (16): Conv1d(32, 64, kernel_size=(5,), stride=(2,), padding=(2,))
        (17): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (18): ELU(alpha=1.0)
        (19): LayerNorm((64, 20), eps=1e-05, elementwise_affine=True)
        (20): Conv1d(64, 128, kernel_size=(5,), stride=(2,), padding=(2,))
        (21): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (22): ELU(alpha=1.0)
        (23): LayerNorm((128, 10), eps=1e-05, elementwise_affine=True)
        (24): AdaptiveAvgPool1d(output_size=1)
        (25): Flatten(start_dim=1, end_dim=-1)
        (26): Linear(in_features=128, out_features=128, bias=True)
      )
    )
    (spatial_pooling): AttentionPooling(
      (transformer_layer): TransformerEncoderLayer(
        (self_attn): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)
        )
        (linear1): Linear(in_features=128, out_features=2048, bias=True)
        (dropout): Dropout(p=0.0, inplace=False)
        (linear2): Linear(in_features=2048, out_features=128, bias=True)
        (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (dropout1): Dropout(p=0.0, inplace=False)
        (dropout2): Dropout(p=0.0, inplace=False)
      )
    )
    (positional_encoding): PositionalEncoding()
    (layer_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
    (transformer_encoder): TransformerEncoder(
      (layers): ModuleList(
        (0-5): 6 x TransformerEncoderLayer(
          (self_attn): MultiheadAttention(
            (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)
          )
          (linear1): Linear(in_features=128, out_features=2048, bias=True)
          (dropout): Dropout(p=0.0, inplace=False)
          (linear2): Linear(in_features=2048, out_features=128, bias=True)
          (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
          (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
          (dropout1): Dropout(p=0.0, inplace=False)
          (dropout2): Dropout(p=0.0, inplace=False)
        )
      )
    )
    (temporal_pooling): AttentionPooling(
      (transformer_layer): TransformerEncoderLayer(
        (self_attn): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)
        )
        (linear1): Linear(in_features=128, out_features=2048, bias=True)
        (dropout): Dropout(p=0.0, inplace=False)
        (linear2): Linear(in_features=2048, out_features=128, bias=True)
        (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (dropout1): Dropout(p=0.0, inplace=False)
        (dropout2): Dropout(p=0.0, inplace=False)
      )
    )
  )
)
```

--------------------------------

### Sleep Staging Model Validation Loop

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

This snippet outlines the validation loop for a sleep staging model. It iterates through the test DataLoader, performs inference, and collects targets, logits, outputs, masks, and file paths. The results are then saved using a `save_data` function. Ensure the model, device, and `tqdm` are imported and initialized.

```python
# Validation loop at the end of each epoch
model.eval()
all_targets = []
all_logits = []
all_outputs = []
all_masks = []
all_paths = []

count = 0
with torch.no_grad():
    for (x_data, y_data, padded_matrix, hdf5_path_list) in tqdm.tqdm(test_loader, desc="Evaluating"):
        x_data, y_data, padded_matrix, hdf5_path_list = x_data.to(device), y_data.to(device), padded_matrix.to(device), list(hdf5_path_list)
        outputs, mask = sleep_staging_model(x_data, padded_matrix)
        all_targets.append(y_data.cpu().numpy())
        all_outputs.append(torch.softmax(outputs, dim=-1).cpu().numpy())
        all_logits.append(outputs.cpu().numpy())
        all_masks.append(mask.cpu().numpy())
        all_paths.append(hdf5_path_list)


save_path = os.path.join(base_save_path, "demo_sleep_staging")
os.makedirs(save_path, exist_ok=True)

targets_path = os.path.join(save_path, "all_targets.pickle")
outputs_path = os.path.join(save_path, "all_outputs.pickle")
logits_path = os.path.join(save_path, "all_logits.pickle")
mask_path = os.path.join(save_path, "all_masks.pickle")
file_paths = os.path.join(save_path, "all_paths.pickle")

save_data(all_targets, targets_path)
save_data(all_outputs, outputs_path)
save_data(all_logits, logits_path)
save_data(all_masks, mask_path)
save_data(all_paths, file_paths)
```

--------------------------------

### Compute and Visualize F1 Scores and Confusion Matrix

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Calculates F1 scores for each class and generates a normalized confusion matrix with percentages. Requires `f1_score` and `confusion_matrix` from scikit-learn, and plotting libraries like matplotlib and seaborn. Ensure `class_labels` are defined correctly.

```python
# Step 1: Get predicted labels (argmax on probabilities)
predicted_labels = np.argmax(all_outputs_filtered, axis=1)

fontsize = 12

# Step 2: Compute F1 score for each class
f1_scores = f1_score(all_targets_filtered, predicted_labels, average=None, labels=range(len(class_labels)))
for idx, label in enumerate(class_labels):
    print(f"F1 Score for {label}: {f1_scores[idx]:.3f}")

# Step 3: Create a confusion matrix and normalize it by row to get percentages
conf_matrix = confusion_matrix(all_targets_filtered, predicted_labels, labels=range(len(class_labels)))
conf_matrix_percent = conf_matrix / conf_matrix.sum(axis=1, keepdims=True) * 100

# Plotting the confusion matrix with percentages
plt.figure(figsize=(6, 4))
sns.heatmap(
    conf_matrix_percent,
    annot=True,
    fmt=".1f",
    cmap="Blues",
    xticklabels=class_labels,
    yticklabels=class_labels,
    annot_kws={\"size\": fontsize},  # Font size for numbers inside the heatmap
    cbar_kws={\"shrink\": 1},  # Adjust colorbar size
)

# Customizing axis labels and ticks
plt.xlabel("Predicted Labels", fontsize=fontsize)
plt.ylabel("True Labels", fontsize=fontsize)
plt.xticks(fontsize=12, ha="center")  # Font size for x-axis tick labels with rotation
plt.yticks(fontsize=12)  # Font size for y-axis tick labels

# Adjust layout and save the figure
plt.tight_layout()
plt.show()
```

--------------------------------

### Load Disease Labels

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads the mapping between disease indices and their corresponding phecodes and phenotypes from a CSV file.

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

--------------------------------

### Detailed Shapes of Model Outputs and Targets

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

This snippet provides the detailed shapes of all collected outputs from the validation loop: logits, softmax outputs, targets, and masks. This is useful for debugging and ensuring data consistency before further processing.

```python
all_logits[0].shape, all_outputs[0].shape, all_targets[0].shape, all_masks[0].shape
```

--------------------------------

### SleepFM Project BibTeX Entry

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/README.md

The BibTeX entry for the SleepFM paper. Use this for citing the project in academic work.

```bibtex
@article{thapa2026multimodal,
  title={A multimodal sleep foundation model for disease prediction},
  author={Thapa, Rahul and Kjaer, Magnus Ruud and He, Bryan and Covert, Ian and Moore IV, Hyatt and Hanif, Umaer and Ganjoo, Gauri and Westover, M Brandon and Jennum, Poul and Brink-Kjaer, Andreas and others},
  journal={Nature Medicine},
  pages={1--11},
  year={2026},
  publisher={Nature Publishing Group US New York}
}
```

--------------------------------

### Map Predictions to Labels

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Integrates the model's prediction outputs, event indicators, and event times into the labels DataFrame for further analysis.

```python
labels_df["output"] = all_outputs[0]
labels_df["is_event"] = all_is_event[0]
labels_df["event_time"] = all_event_times[0]
```

--------------------------------

### Load and Prepare Clinical Data for Disease Prediction

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Loads disease labels from CSV files and resolves HDF5 paths for a given split. Filters studies based on available labels and optionally truncates the dataset. Builds an index map linking HDF5 file paths to their corresponding labels.

```python
is_event_df = pd.read_csv(os.path.join(self.config["labels_path"], "is_event.csv"))
event_time_df = pd.read_csv(os.path.join(self.config["labels_path"], "time_to_event.csv"))

is_event_df = is_event_df.set_index('Study ID')
event_time_df = event_time_df.set_index('Study ID')

# --- Resolve HDF5 paths (explicit precedence) ---
if hdf5_paths:
    # Use provided paths directly
    hdf5_paths = [f for f in hdf5_paths if os.path.exists(f)]
else:
    # Load from split file
    split_paths = load_data(config["split_path"])[split]
    hdf5_paths = [f for f in split_paths if os.path.exists(f)]

# Filter by available demo labels
hdf5_paths = [
    f for f in hdf5_paths
    if os.path.basename(f).split(".")[0] in study_ids
]

# Optional truncation
if config.get("max_files"):
    hdf5_paths = hdf5_paths[:config["max_files"]]

labels_dict = {}
# Loop over each study_id
for study_id in tqdm.tqdm(study_ids):
    # Extract the row as a whole for both dataframes (faster than iterating over columns)
    is_event_row = list(is_event_df.loc[study_id].values)
    event_time_row = list(event_time_df.loc[study_id].values)
    demo_feats = list(demo_labels_df.loc[study_id].values)

    # values = [[event_time, is_event] for is_event, event_time in zip(is_event_row, event_time_row)]
    labels_dict[study_id] = {
        "is_event": is_event_row,
        "event_time": event_time_row, 
        "demo_feats": demo_feats
    }

# --- Build index map ---
self.index_map = [
    (path, labels_dict[os.path.basename(path).split(".")[0]])
    for path in hdf5_paths
]

print(f"Number of files in {split} set: {len(hdf5_paths)}")
print(f"Number of files to be processed in {split} set: {len(self.index_map)}")

self.total_len = len(self.index_map)
self.max_seq_len = config["model_params"]["max_seq_length"]

if self.total_len == 0:
    raise ValueError(f"No valid HDF5 files found for split='{split}'.")
```

--------------------------------

### Model Evaluation Loop

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Evaluates the trained model on the test dataset. It iterates through the DataLoader, performs inference, and collects outputs, event times, and event indicators.

```python
model.eval()
all_event_times = []
all_is_event = []
all_outputs = []
all_paths = []

with torch.no_grad():
    for item in tqdm.tqdm(test_loader, desc="Evaluating"):
        x_data, event_times, is_event, demo_feats, padded_matrix, hdf5_path_list = item
        x_data, event_times, is_event, demo_feats, padded_matrix, hdf5_path_list = x_data.to(device), event_times.to(device), is_event.to(device), demo_feats.to(device), padded_matrix.to(device), list(hdf5_path_list)
        outputs = model(x_data, padded_matrix, demo_feats)
    
        logits = outputs.cpu().numpy()
        all_outputs.append(logits)
        all_event_times.append(event_times.cpu().numpy())
        all_is_event.append(is_event.cpu().numpy())
        all_paths.append(hdf5_path_list)

all_outputs = np.concatenate(all_outputs, axis=0)
all_event_times = np.concatenate(all_event_times, axis=0)
all_is_event = np.concatenate(all_is_event, axis=0)
all_paths = np.concatenate(all_paths)

outputs_path = os.path.join(save_path, "all_outputs.pickle")
event_times_path = os.path.join(save_path, "all_event_times.pickle")
is_event_path = os.path.join(save_path, "all_is_event.pickle")
file_paths = os.path.join(save_path, "all_paths.pickle")

save_data(all_outputs, outputs_path)
save_data(all_event_times, event_times_path)
save_data(all_is_event, is_event_path)
save_data(all_paths, file_paths)
```

--------------------------------

### Shape of Model Outputs

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

This snippet displays the shapes of the collected model outputs and targets after the validation loop. It confirms the dimensions for logits, outputs, targets, and masks, which are crucial for subsequent analysis or evaluation.

```python
all_outputs[0].shape, all_targets[0].shape
```

--------------------------------

### Custom Sleep Event Classification Dataset

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

A custom PyTorch Dataset class for loading and preprocessing sleep event classification data from HDF5 files and label CSVs. It handles data indexing, context windowing, and channel selection.

```python
class SleepEventClassificationDataset(Dataset):
    def __init__(
        self,
        config,
        channel_groups,
        hdf5_paths,
        label_files,
        split="train",
    ):
        self.config = config
        self.max_channels = self.config["max_channels"]
        self.context = int(self.config["context"])
        self.channel_like = self.config["channel_like"]

        self.max_seq_len = config["model_params"]["max_seq_length"]

        # --- Build label lookup: {study_id: label_csv_path} ---
        # study_id = filename without extension, e.g. "SSC_12345"
        labels_dict = {
            os.path.basename(p).rsplit(".", 1)[0]: p
            for p in label_files
            if os.path.exists(p)
        }

        # --- Filter to HDF5s that exist and have a matching label file ---
        hdf5_paths = [p for p in hdf5_paths if os.path.exists(p)]
        hdf5_paths = [
            p for p in hdf5_paths
            if os.path.basename(p).rsplit(".", 1)[0] in labels_dict
        ]

        if config.get("max_files"):
            hdf5_paths = hdf5_paths[: config["max_files"]]

        self.hdf5_paths = hdf5_paths
        self.labels_dict = labels_dict

        # --- Build index map ---
        # Each item is (hdf5_path, label_path, start_index)
        if self.context == -1:
            self.index_map = [
                (p, labels_dict[os.path.basename(p).rsplit(".", 1)[0]], -1)
                for p in self.hdf5_paths
            ]
        else:
            self.index_map = []
            loop = tqdm(self.hdf5_paths, total=len(self.hdf5_paths), desc=f"Indexing {split} data")
            for hdf5_file_path in loop:
                file_prefix = os.path.basename(hdf5_file_path).rsplit(".", 1)[0]
                label_path = labels_dict[file_prefix]

                with h5py.File(hdf5_file_path, "r") as hf:
                    dset_names = list(hf.keys())
                    if len(dset_names) == 0:
                        continue

                    # Use first dataset to define length (same as your original behavior)
                    first_name = dset_names[0]
                    dataset_length = hf[first_name].shape[0]

                for i in range(0, dataset_length, self.context):
                    self.index_map.append((hdf5_file_path, label_path, i))

        # If you have logger, keep; otherwise you can remove these.
        # logger.info(f"Number of files in {split} set: {len(self.hdf5_paths)}")
        # logger.info(f"Number of files to be processed in {split} set: {len(self.index_map)}")

        self.total_len = len(self.index_map)

    def __len__(self):
        return self.total_len

    def get_index_map(self):
        return self.index_map

    def __getitem__(self, idx):
        hdf5_path, label_path, start_index = self.index_map[idx]

        labels_df = pd.read_csv(label_path)
        labels_df["StageNumber"] = labels_df["StageNumber"].replace(-1, 0)

        y_data = labels_df["StageNumber"].to_numpy()
        if self.context != -1:
            y_data = y_data[start_index : start_index + self.context]

        x_data = []
        with h5py.File(hdf5_path, "r") as hf:
            dset_names = list(hf.keys())

            for dataset_name in dset_names:
                if dataset_name in self.channel_like:
                    if self.context == -1:
                        x_data.append(hf[dataset_name][:])
                    else:
                        x_data.append(hf[dataset_name][start_index : start_index + self.context])

        if not x_data:
            # Skip this data point if x_data is empty
            return self.__getitem__((idx + 1) % self.total_len)

        x_data = np.array(x_data)  # (C, T, F) assuming each channel returns (T, F)
        x_data = torch.tensor(x_data, dtype=torch.float32)
        y_data = torch.tensor(y_data, dtype=torch.float32)

        min_length = min(x_data.shape[1], len(y_data))
        x_data = x_data[:, :min_length, :]
        y_data = y_data[:min_length]

        return x_data, y_data, self.max_channels, self.max_seq_len, hdf5_path
```

--------------------------------

### Extract and Save Embeddings from Clinical Data

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Processes batches of clinical data to extract and save both 5-second level and 5-minute aggregated embeddings into HDF5 files. Requires model, dataloader, device, and configuration to be set up.

```python
with torch.no_grad():
    with tqdm.tqdm(total=len(dataloader)) as pbar:
        for batch in dataloader:
            batch_data, mask_list, file_paths, dset_names_list, chunk_starts = batch
            (bas, resp, ekg, emg) = batch_data
            (mask_bas, mask_resp, mask_ekg, mask_emg) = mask_list

            bas = bas.to(device, dtype=torch.float)
            resp = resp.to(device, dtype=torch.float)
            ekg = ekg.to(device, dtype=torch.float)
            emg = emg.to(device, dtype=torch.float)

            mask_bas = mask_bas.to(device, dtype=torch.bool)
            mask_resp = mask_resp.to(device, dtype=torch.bool)
            mask_ekg = mask_ekg.to(device, dtype=torch.bool)
            mask_emg = mask_emg.to(device, dtype=torch.bool)

            embeddings = [
                model(bas, mask_bas),
                model(resp, mask_resp),
                model(ekg, mask_ekg),
                model(emg, mask_emg),
            ]

            # Model gives two kinds of embeddings. Granular 5 second-level embeddings and aggregated 5 minute-level embeddings. We save both of them below. 

            embeddings_new = [e[0].unsqueeze(1) for e in embeddings]

            for i in range(len(file_paths)):
                file_path = file_paths[i]
                chunk_start = chunk_starts[i]
                subject_id = os.path.basename(file_path).split('.')[0]
                output_path = os.path.join(output_5min_agg, f"{subject_id}.hdf5")

                with h5py.File(output_path, 'a') as hdf5_file:
                    for modality_idx, modality_type in enumerate(config["modality_types"]):
                        if modality_type in hdf5_file:
                            dset = hdf5_file[modality_type]
                            chunk_start_correct = chunk_start // (embed_dim * 5 * 60)
                            chunk_end = chunk_start_correct + embeddings_new[modality_idx][i].shape[0]
                            if dset.shape[0] < chunk_end:
                                dset.resize((chunk_end,) + embeddings_new[modality_idx][i].shape[1:])
                            dset[chunk_start_correct:chunk_end] = embeddings_new[modality_idx][i].cpu().numpy()
                        else:
                            hdf5_file.create_dataset(modality_type, data=embeddings_new[modality_idx][i].cpu().numpy(), chunks=(embed_dim,) + embeddings_new[modality_idx][i].shape[1:], maxshape=(None,) + embeddings_new[modality_idx][i].shape[1:])

            embeddings_new = [e[1] for e in embeddings]

            for i in range(len(file_paths)):
                file_path = file_paths[i]
                chunk_start = chunk_starts[i]
                subject_id = os.path.basename(file_path).split('.')[0]
                output_path = os.path.join(output, f"{subject_id}.hdf5")

                with h5py.File(output_path, 'a') as hdf5_file:
                    for modality_idx, modality_type in enumerate(config["modality_types"]):
                        if modality_type in hdf5_file:
                            dset = hdf5_file[modality_type]
                            chunk_start_correct = chunk_start // (embed_dim * 5)
                            chunk_end = chunk_start_correct + embeddings_new[modality_idx][i].shape[0]
                            if dset.shape[0] < chunk_end:
                                dset.resize((chunk_end,) + embeddings_new[modality_idx][i].shape[1:])
                            dset[chunk_start_correct:chunk_end] = embeddings_new[modality_idx][i].cpu().numpy()
                        else:
                            hdf5_file.create_dataset(modality_type, data=embeddings_new[modality_idx][i].cpu().numpy(), chunks=(embed_dim,) + embeddings_new[modality_idx][i].shape[1:], maxshape=(None,) + embeddings_new[modality_idx][i].shape[1:])
            pbar.update()
```

--------------------------------

### Define Custom Dataset Class

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Defines a custom PyTorch Dataset class for loading and preprocessing data for disease prediction. It handles loading demographic features from a CSV file and setting the 'Study ID' as the index.

```python
class DiagnosisFinetuneFullCOXPHWithDemoDataset(Dataset):
    def __init__(self, 
                 config,
                 channel_groups,
                 hdf5_paths=None,
                 demo_labels_path=None,
                 split="train"):

        self.config = config
        self.channel_groups = channel_groups
        self.max_channels = self.config["max_channels"]

        # --- Load demographic features ---
        if not demo_labels_path:
            demo_labels_path = config["demo_labels_path"]

        demo_labels_df = pd.read_csv(demo_labels_path)
        demo_labels_df = demo_labels_df.set_index("Study ID")
        study_ids = set(demo_labels_df.index)
```

--------------------------------

### Inspect Prediction Shapes

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Checks the shapes of the collected prediction outputs, event times, and event indicators to ensure data integrity after evaluation.

```python
all_outputs.shape, all_event_times.shape, all_is_event.shape
```

--------------------------------

### Filter Data Using Mask

Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb

Applies a binary mask to filter specific elements from flattened data arrays. The mask should have a boolean dtype.

```python
mask_filter = all_masks_flat == 0

# Apply the mask to each flattened array
all_logits_filtered = all_logits_flat[mask_filter]
all_outputs_filtered = all_outputs_flat[mask_filter]
all_targets_filtered = all_targets_flat[mask_filter]
```
Tokens:
Raw