HEART: Learning better representation of EHR data with a heterogeneous relation-aware transformer

This is our PyTorch implementation for the paper:

Tinglin Huang, Syed Asad Rizvi, Rohan Krishna Thakur, Vimig Socrates, Meili Gupta, David van Dijk, R. Andrew Taylor, Rex Ying (2024). HEART: Learning better representation of EHR data with a heterogeneous relation-aware transformer. Paper Link. In Journal of Biomedical Informatics 159 (2024): 104741.

Dataset Preparation

We have provided the preprocessing scripts in dataset/ for MIMIC-III and eICU datasets respectively. Please first download the datasets from the following link:

• MIMIC-III: https://mimic.mit.edu/docs/iii/
• eICU: https://eicu-crd.mit.edu/about/eicu/

Requirements

The code has been tested running under Python 3.10.14. The required packages are as follows:

• pytorch == 2.3.0
• torch_geometric == 2.5.3
• torch_scatter == 2.1.2
• einops == 0.8.0
• numpy
• pandas
• scikit-learn
• tqdm
• wandb

Once you finished these installation, please run install the package by running:

pip install -e .

Organization

The code is organized as follows:

• app/: the main code for training and testing the model
  • finetune.py: the pipeline for finetuning the model on downstream tasks
  • pretrain.py: the pipeline for pretraining the model on the pretraining task
• dataset/: the code for data processing
  • eICU.ipynb: dataset preprocessing for eICU
  • MIMIC-III.ipynb: dataset preprocessing for MIMIC-III
• models/
  • gnn.py: implementation of the graph attention for the encounter-level attention
  • HEART.py: implementation of the pretraining and finetuning model
  • transformer_rel.py: implementation of the transformer with heterogeneous relations
  • transformer.py: implementation of the transformer
• utils/: utility functions including data loading pipeline

Usage

Note that the data paths and the checkpoint paths in app/pretrain.py and app/finetune.py are hard-coded (/home/username/ehr_bert/...), please update them to your own paths before running. Then run:

python heart/app/pretrain.py --dataset mimic
python heart/app/finetune.py --dataset mimic --task death