DeepHealthExpo/README.md

DeepHealthExpo

Next-token DeepHealth training code for disease-event sequence modeling with optional extra/exposure information.

This repository is a clean code-only extraction from the main DeepHealth project. It keeps the next-token training path and reusable model/data utilities, while excluding large UKB data files, trained checkpoints, result folders, and all-future training entry points.

Included

• train_next_step.py: next-token / UTS training entry point.
• dataset.py: next-step event sequence dataset with unified extra-info tokens.
• models.py, backbones.py: DeepHealth Transformer backbone.
• losses.py, readouts.py, targets.py: training targets, losses, and readout utilities.
• evaluate_auc.py, evaluate_token_auc.py: next-token checkpoint evaluation utilities.
• prepare_data.py, prepare_event_dates.py, event_date_utils.py: data preparation helpers.
• extra_info_types_*.txt: reusable extra-info type selections.

Not Included

The repository intentionally does not include raw or derived UKB arrays, split files, checkpoints, or run outputs.

Expected local data files for training normally include:

ukb_event_data.npy
ukb_other_info.npy
ukb_basic_info.csv
ukb_train_eid.csv
ukb_val_eid.csv
ukb_test_eid.csv
cate_types.csv

labels.csv and field_ids_enriched.csv are included because they define the model vocabulary and preparation metadata.

Example

python train_next_step.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --extra_info_types_file extra_info_types_exposure_only.txt \
  --target_mode uts \
  --time_mode relative

For strict next-token Delphi-style training:

python train_next_step.py --target_mode delphi2m --readout_name token

Exposure Modeling Direction

For onset-aligned environmental exposure parquet files, the first intended extension is single-stream event enhancement:

disease event token + pre-onset exposure embedding -> same next-token Transformer

The key constraint is that a disease event's own pre-onset exposure must not be used to predict that same disease event.

Pretrain the exposure encoder as a denoising autoencoder using training-set EIDs:

python train_exposure_autoencoder.py \
  --exposure_cache_dir ukb_exposure_cache \
  --train_eid_file ukb_train_eid.csv \
  --val_eid_file ukb_val_eid.csv

The best checkpoint contains the encoder and normalization statistics needed to generate fixed exposure embeddings. Multi-GPU pretraining follows the main trainer interface: add --data_parallel --gpu_ids 0,1,2,3. For efficient multi-GPU training, launch one process per GPU with DDP:

torchrun --standalone --nproc_per_node=4 train_exposure_autoencoder.py \
  --exposure_cache_dir ukb_exposure_cache \
  --batch_size 128

In DDP mode, --batch_size is the global batch size and must be divisible by the number of processes.

The trainer also writes last.pt after every epoch so interrupted runs can be continued. Pass the run directory to reuse the original train_config.json; the trainer will load last.pt when available and fall back to best.pt for older runs:

python train_exposure_autoencoder.py \
  --resume_checkpoint runs/exposure_ae_RUN

Encode every cached exposure window once:

python encode_exposure_cache.py \
  --exposure_cache_dir ukb_exposure_cache \
  --checkpoint runs/exposure_ae_RUN/best.pt

The next-step trainer reads exposure_embeddings.npy directly and does not run TimesNet:

torchrun --standalone --nproc_per_node=4 train_next_step.py \
  --exposure_cache_dir ukb_exposure_cache \
  --batch_size 128

Training-channel statistics are cached at <exposure_cache_dir>/train_channel_stats.npz; use --recompute_channel_stats only when a forced refresh is needed.