MEDS-EIC-AR 0.1.11

A simple auto-regressive, 'everything-is-code' style model for MEDS datasets

MEDS "Everything-is-code" Autoregressive Model

A MEDS, "Everything-is-code" style Autoregressive Generative Model, capable of zero-shot inference.

This is based on the MEDS-Torch model of the same name.

Installation

pip install MEDS-EIC-AR

Optional Dependencies

WandB

If you want to use WandB for logging, you can install it via:

pip install MEDS-EIC-AR[wandb]

MLFlow

If you want to use MLFlow for logging, you can install it via:

pip install MEDS-EIC-AR[mlflow]

This will also install psutil and pynvml as dependencies, to enable MLFlow tracking of system CPU and GPU resources, which is enabled by default or can be controlled via the MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING environment variable. See the MLFlow documentation for more details.

Flash Attention

For using flash attention, you need to subsequently install flash attention as well. This can often be done via:

pip install flash-attn --no-build-isolation

If you encounter errors, see the flash-attn package documentation.

Usage

1. Pre-process your data

You have three directories:

1. $RAW_MEDS_DIR -- The raw MEDS data directory that you want to pre-process.
2. $INTERMEDIATE_DIR -- An intermediate directory where the partially processed data will be stored prior to tokenization and tensorization.
3. $FINAL_DATA_DIR -- The final output directory where the tokenized and tensorized data will be stored. This directory is suitable for use in loading the data with meds-torch-data.

Run:

MEICAR_process_data input_dir="$RAW_MEDS_DIR" \
    intermediate_dir="$INTERMEDIATE_DIR" \
    output_dir="$FINAL_DATA_DIR"

[!NOTE] If your data is not sharded by split at the outset, you will need to add the do_reshard=True command line parameter to the MEICAR_process_data command, which ensures the system reshards the data to be sub-sharded by split before beginning pre-processing.

You can also run this in demo mode, which lowers the filtering thresholds significantly so the script does not filter out all data:

MEICAR_process_data ... do_demo=True

You can exert more fine-grained control on the filtering with the following environment variables:

1. MIN_SUBJECTS_PER_CODE: How many subjects must a given code be observed within to be included in the final vocabulary? Note that this excludes some sentinel codes which are always retained.
2. MIN_EVENTS_PER_SUBJECT: How many events must a subject have to be included in the final dataset?

2. Pre-train the model

You can pre-train the model using the MEICAR_pretrain command. To use this, let us assume you have a new directory to store the pretrained model artifacts called $PRETRAINED_MODEL_DIR. Then, you can run:

MEICAR_pretrain datamodule.config.tensorized_cohort_dir="$FINAL_DATA_DIR" \
    output_dir="$PRETRAINED_MODEL_DIR" \
    datamodule.batch_size=32

to train the model for 10 epochs.

This uses a Hydra configuration system, with the root config located in the _pretrain.yaml file. You can override any of the nested configuration parameters (as shown above via datamodule.config.tensorized_cohort_dir on the command line, though you will more likely materialize an experimental configuration file to disk in yaml form and overwrite the config path and name directly in the normal hydra manner.

[!WARNING] Tests here only validate that the model runs without errors and (in demo mode) runs without producing nans or invalid values. It has not yet been assessed to ensure it runs to convergence, etc.

3. Zero-shot Inference

Zero-shot inference consists of two steps:

1. Given a task cohort and a pre-trained model, for each sample in the task cohort, generate future trajectories from those inputs forward with the pre-trained model and save them to disk in a pseudo-MEDS format.
2. Resolve these generated trajectories into concrete, probabilistic predictions for the task cohort.

3.1 Generate Trajectories for a task spec.

You can directly generate trajectories using the MEICAR_generate_trajectories command. This requires a few more configuration parameters than the pre-training step, so let's go through those:

1. You need to specify the task labels directory in the datamodule.config.task_labels_dir parameter.
2. You need to specify the model initialization directory in the model_initialization_dir parameter. This is the output directory of the pre-train step.
3. You need to specify how you want to trade-off between allowed input context size and the maximum possible generated trajectory length. The former allows you to use more of the patient's record, but the latter controls how far into the future you can predict. This can be configured with one of three parameters in the seq_lens part of the config. If you set:
   • seq_lens.generation_context_size, that will be the maximum length of the input context, and the remaining length of the pretrained model's maximum sequence length will be used for generation.
   • seq_lens.max_generated_trajectory_len, that will be the maximum length of the generated trajectory, and the remaining length of the pretrained model's maximum sequence length will be used for the input.
   • seq_lens.frac_seq_len_as_context, that will be the fraction of the pretrained model's maximum sequence length that will be used for the input context, and the remaining length will be used for generation. This is set by default to 0.25, which means that 25% of the maximum sequence length will be used for the input context, and 75% will be used for generation. If you wish to use another mode on the command line, be sure to set this to null to disable it.
4. Lastly, you need to specify how many trajectories per task sample you wish to generate, and for which splits you wish to generate samples. You can do this via the inference.generate_for_splits and inference.N_trajectories_per_task_sample parameters. The former is a list of splits to generate and the latter is the number of trajectories to generate per task sample. The default is to generate 20 trajectories for each task sample in the tuning and held out splits.

After these are set, you can run the following command to generate trajectories for a task cohort:

MEICAR_generate_trajectories \
    output_dir="$GENERATED_TRAJECTORIES_DIR" \
    model_initialization_dir="$PRETRAINED_MODEL_DIR" \
    datamodule.config.tensorized_cohort_dir="$FINAL_DATA_DIR" \
    datamodule.config.task_labels_dir="$TASK_ROOT_DIR/$TASK_NAME" \
    datamodule.batch_size=32

This will generate trajectories for the task cohort and save them in the format: $GENERATED_TRAJECTORIES_DIR/$SPLIT/$SAMPLE.parquet.

See the documentation for format_trajectories for more details on the format of the generated trajectories.

[!WARNING] The tests here only validate that this runs without errors and produces trajectory files that are valid, non-identical across different samples, and containing the right subjects. It has not yet been assessed to ensure full correctness.

[!NOTE] The generated trajectories from this model are saved in the schema defined in the MEDS_trajectory_evaluation.schema.GeneratedTrajectorySchema format, and can be used with that package's evaluation tools.

3.2 Resolve Trajectories into Predictions.

Not yet implemented.

Documentation

Configuration and Controlling Model Structure

This model is configured via Hydra and PyTorch lightning. The configuration structure of this repository is as follows:

>>> print_directory("./src/MEDS_EIC_AR/configs", config=PrintConfig(file_extension=".yaml"))
├── _demo_generate_trajectories.yaml
├── _demo_pretrain.yaml
├── _generate_trajectories.yaml
├── _pretrain.yaml
├── datamodule
│   ├── default.yaml
│   ├── generate_trajectories.yaml
│   └── pretrain.yaml
├── inference
│   ├── default.yaml
│   └── demo.yaml
├── lightning_module
│   ├── LR_scheduler
│   │   └── get_cosine_schedule_with_warmup.yaml
│   ├── default.yaml
│   ├── demo.yaml
│   ├── large.yaml
│   ├── medium.yaml
│   ├── metrics
│   │   └── default.yaml
│   ├── micro.yaml
│   ├── model
│   │   ├── default.yaml
│   │   ├── demo.yaml
│   │   ├── large.yaml
│   │   ├── medium.yaml
│   │   ├── micro.yaml
│   │   └── small.yaml
│   ├── optimizer
│   │   └── adamw.yaml
│   └── small.yaml
└── trainer
    ├── callbacks
    │   ├── default.yaml
    │   ├── early_stopping.yaml
    │   ├── learning_rate_monitor.yaml
    │   └── model_checkpoint.yaml
    ├── default.yaml
    ├── demo.yaml
    └── logger
        ├── csv.yaml
        ├── mlflow.yaml
        └── wandb.yaml

Logging with wandb

You can activate the wandb logger by overriding the trainer logger to wandb:

MEICAR_pretrain trainer.logger=wandb

The configuration file configs/trainer/logger/wandb.yaml exposes a tags field. Hydra makes the selected configuration groups available via hydra.runtime.choices. These can be referenced to automatically tag the run. For example:

tags:
  - ${hydra:runtime.choices.lightning_module/model}

This automatically tags each run with the selected model size (e.g. small, medium, large). Hydra currently cannot append to a list with a default value. To add your own tags you must override the list and include the default tag yourself:

MEICAR_pretrain trainer.logger=wandb \
  trainer.logger.tags="[${hydra:runtime.choices.lightning_module/model},experiment-1]"

This results in the tags [model_size, "experiment-1"] being sent to wandb.

Output Files

The output files of the pre-training step are stored in the directory specified by the output_dir parameter and take the following structure:

>>> print_directory(pretrained_model)
├── .logs
│   ├── .hydra
│   │   ├── config.yaml
│   │   ├── hydra.yaml
│   │   └── overrides.yaml
│   └── __main__.log
├── best_model.ckpt
├── checkpoints
│   ├── epoch=0-step=1.ckpt
│   ├── epoch=0-step=2.ckpt
│   ├── epoch=1-step=3.ckpt
│   ├── epoch=1-step=4.ckpt
│   └── last.ckpt
├── config.yaml
├── loggers
│   └── csv
│       └── version_0
│           ├── hparams.yaml
│           └── metrics.csv
└── resolved_config.yaml