G-Bert
Pre-training of Graph Augmented Transformers for Medication Recommendation
Intro
G-Bert combined the power of Graph Neural Networks and BERT (Bidirectional Encoder Representations from Transformers) for medical code representation and medication recommendation. We use the graph neural networks (GNNs) to represent the structure information of medical codes from a medical ontology. Then we integrate the GNN representation into a transformer-based visit encoder and pre-train it on single-visit EHR data. The pre-trained visit encoder and representation can be fine-tuned for downstream medical prediction tasks. Our model is the first to bring the language model pre-training schema into the healthcare domain and it achieved state-of-the-art performance on the medication recommendation task.
Requirements
- pytorch>=0.4
- python>=3.5
- torch_geometric==1.0.3
Guide
We list the structure of this repo as follows:
.
├── [4.0K] code/
│ ├── [ 13K] bert_models.py % transformer models
│ ├── [5.9K] build_tree.py % build ontology
│ ├── [4.3K] config.py % hyperparameters for G-Bert
│ ├── [ 11K] graph_models.py % GAT models
│ ├── [ 0] __init__.py
│ ├── [9.8K] predictive_models.py % G-Bert models
│ ├── [ 721] run_alternative.sh % script to train G-Bert
│ ├── [ 19K] run_gbert.py % fine tune G-Bert
│ ├── [ 19K] run_gbert_side.py
│ ├── [ 18K] run_pretraining.py % pre-train G-Bert
│ ├── [4.4K] run_tsne.py # output % save embedding for tsne visualization
│ └── [4.7K] utils.py
├── [4.0K] data/
│ ├── [4.9M] data-multi-side.pkl
│ ├── [3.6M] data-multi-visit.pkl % patients data with multi-visit
│ ├── [4.3M] data-single-visit.pkl % patients data with singe-visit
│ ├── [ 11K] dx-vocab-multi.txt % diagnosis codes vocabulary in multi-visit data
│ ├── [ 11K] dx-vocab.txt % diagnosis codes vocabulary in all data
│ ├── [ 29K] EDA.ipynb % jupyter version to preprocess data
│ ├── [ 18K] EDA.py % python version to preprocess data
│ ├── [6.2K] eval-id.txt % validation data ids
│ ├── [6.9K] px-vocab-multi.txt % procedure codes vocabulary in multi-visit data
│ ├── [ 725] rx-vocab-multi.txt % medication codes vocabulary in multi-visit data
│ ├── [2.6K] rx-vocab.txt % medication codes vocabulary in all data
│ ├── [6.2K] test-id.txt % test data ids
│ └── [ 23K] train-id.txt % train data ids
└── [4.0K] saved/
└── [4.0K] GBert-predict/ % model files to reproduce our result
├── [ 371] bert_config.json
└── [ 12M] pytorch_model.bin
Preprocessing Data
We have released the preprocessing codes named data/EDA.ipynb to process data using raw files from MIMIC-III dataset. You can download data files from MIMIC and get necessary mapping files from GAMENet.
Quick Test
To validate the performance of G-Bert, you can run the following script since we have provided the trained model binary file and well-preprocessed data.
cd code/
python run_gbert.py --model_name GBert-predict --use_pretrain --pretrain_dir ../saved/GBert-predict --graph
Cite
Please cite our paper if you find this code helpful:
@article{shang2019pre,
title={Pre-training of Graph Augmented Transformers for Medication Recommendation},
author={Shang, Junyuan and Ma, Tengfei and Xiao, Cao and Sun, Jimeng},
journal={arXiv preprint arXiv:1906.00346},
year={2019}
}
Acknowledgement
Many thanks to the open source repositories and libraries to speed up our coding progress.
167 Jan 08, 2023
93 Dec 08, 2022
7 Dec 23, 2022
12 Dec 23, 2022
3.1k Dec 28, 2022
69 Dec 12, 2022
95 Jan 03, 2023
3 Aug 17, 2022
48 Nov 29, 2022
0 Sep 28, 2022
13 Jan 02, 2023
3 Jun 03, 2022
11 Aug 29, 2022
30 Dec 07, 2022
16 Nov 07, 2022
2 Dec 12, 2022
75 Dec 02, 2022
331 Jan 03, 2023
9 Oct 30, 2022
0 Feb 05, 2022