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 both model_state_dict, an encoder_state_dict compatible with the default gated TimesNetExposureEncoder, and the channel normalization statistics needed when the encoder is attached to DeepHealth. 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 end-to-end next-step trainer supports the same DDP launch pattern:

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

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