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:

```text
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

```bash
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:

```bash
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:

```text
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:

```bash
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:

```bash
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:

```bash
python train_exposure_autoencoder.py \
  --resume_checkpoint runs/exposure_ae_RUN
```

Encode every cached exposure window once:

```bash
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:

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

To train the exposure-sequence ablation, which removes disease-token
embeddings while retaining sex and age encodings, add:

```bash
python train_next_step.py \
  --exposure_cache_dir ukb_exposure_cache \
  --input_ablation exposure_only
```
Training-channel statistics are cached at
`<exposure_cache_dir>/train_channel_stats.npz`; use
`--recompute_channel_stats` only when a forced refresh is needed.