April 25th 2022
This is the code repository for the final project of UIUC's CS598 Deep Learning for Healthcare.
This is the reproduced code for the paper DeepNote-GNN: predicting hospital readmission using clinical notes and patient network (link) by Sara Nouri Golmaei and Xiao Luo. The proposed DeepNote model leverages the large-scale pretrained ClinicalBERT (link) and Graph Convolutional Network (GCN) on the constructed patient network to jointly model readmission rate.
This repository implemented the DeepNote model as well as all the baseline models (finetuning ClinicalBERT, bag-of-word, BiLSTM with word2vec) mentioned in the paper. We implement the models with PyTorch, PyTorch-Geometric for graph neural networks, and HuggingFace's Transformers for the pretrained ClinicalBERT.
The performance on the discharge
subset is shown in the below table.
Model | AUROC | AURPC | R@P80% |
---|---|---|---|
DeepNote | 0.8596 | 0.8520 | 0.6759 |
DeepNote-GAT | 0.8591 | 0.8403 | 0.6793 |
ClinicalBert | 0.7721 | 0.7730 | 0.2899 |
BiLSTM | 0.5318 | 0.5778 | 0.0942 |
Bag-of-Word | 0.6899 | 0.7147 | 0.1744 |
The performance on the 3day
subset is shown in the below table.
Model | AUROC | AURPC | R@P80% |
---|---|---|---|
DeepNote | 0.6345 | 0.6140 | 0.0380 |
DeepNote-GAT | 0.6338 | 0.6211 | 0.0342 |
ClinicalBert | 0.6066 | 0.6217 | 0.2421 |
BiLSTM | 0.5054 | 0.5359 | 0.1159 |
Bag-of-Word | 0.6503 | 0.6688 | 0.1377 |
Install the required packages using the following command:
pip install -r requirements.txt
The paper used the MIMIC-III dataset, a large, freely-available database comprising de-identified health-related data. The dataset cannot be redistributed, so you should apply for the access to the dataset. Only the ADMISSIONS.csv
and NOTEEVENTS.csv
are needed, which are 1.1GB total in size.
The data are then preprocessed to extract the readmission information. We adopted the script from ClinicalBERT's repo, as was done by the authors of the paper. Run the preprocessing script with the following command:
python preprocess.py
The script will generate two subset, called discharge
and 3day
respectively under the data_root
folder. The usage is shown below.
usage: preprocessing script
optional arguments:
-h, --help show this help message and exit
--data-root DATA_ROOT (default: './data/')
--seed SEED (default: 1)
Then, use the pretrained ClinicalBERT to extract the representation for each clinical note. The pretrained weights are available at HuggingFace's model hub (link). However, the config file is missing the model_type
key, which will cause an error if directly loaded from the model hub. So you need to download the model and add "model_type": "bert"
to config.json
. After that, run the pretrained model using:
python run_pretrained.py
The script will generate train.pt
, val.pt
, and test.pt
under the data folder. The usage is shown below:
usage: Running pretrained clinicalBERT
optional arguments:
-h, --help show this help message and exit
--data-root DATA_ROOT (default: './data/discharge/')
--model-dir MODEL_DIR (default: './checkpoints/clinicalbert/')
--batch-size BATCH_SIZE (default: 128)
--max-length MAX_LENGTH (default: 512)
Using train.py
to train the DeepNote model as well as the baseline models. For example, you can run the following script to train the DeepNote model:
python train.py configs/deepnote.yml --savename deepnote
The usage is shown below:
positional arguments:
config
optional arguments:
-h, --help show this help message and exit
--device model device (default: 'cuda')
--logdir log folder for tensorboard and saving the model (default: './logs')
--savename log name for tensorboard (default: 'test')
--resume resume checkpoint (default: None)
The configs
folder also includes the config file for finetuning ClinicalBERT and training BiLSTM model and the GAT version of DeepNote. Most of the arguments in the config file is self-explanatory. You can modify them to try different hyperparameter settings.
To run BiLSTM with word2vec, you should first extract the word2vec weights with gensim by running the following script:
pip install gensim
python gen_word2vec.py
The usage is shown below:
usage: Generate word2vec embeddings
optional arguments:
-h, --help show this help message and exit
--save-path SAVE_PATH (default: './checkpoints/word2vec/')
To run the bag-of-word model, use the following script:
python run_bow.py
The usage is shown below:
usage: Running Bag-of-Word model with logistic regression
optional arguments:
-h, --help show this help message and exit
--data-root DATA_ROOT (default: './data/discharge/')
--max-feat MAX_FEAT (default: 5000)
We also provide the code for testing the impact of different thresholds for building graph on the final performance. Run the following command:
python run_diff_threshold.py
The script will generate a CSV file diff_thres.csv
with the following fields:
['threshold', 'num_edges', 'sparsity', 'test_loss', 'acc', 'roc', 'prc', 'rp80']
The usage is shown below:
usage: Test different thresholds for building graph
positional arguments:
config
optional arguments:
-h, --help show this help message and exit
--thres-begin THRES_BEGIN (default: 0.)
--thres-end THRES_END (default: 0.999)
--num-thres NUM_THRES (default: 50)
--device DEVICE (default: 'cuda')