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.
96 Dec 13, 2022
25 Dec 08, 2022
6 Feb 24, 2022
55 Dec 28, 2022
10 Jun 13, 2022
0 May 04, 2022
1k Jan 04, 2023
1k Dec 30, 2022
51 Jan 01, 2023
19 Nov 28, 2022
7 Aug 31, 2022
62 Dec 21, 2022
36 Dec 12, 2022
3 Sep 16, 2022
6 Jan 24, 2022
78 Dec 10, 2022
23 Dec 08, 2022
845 Nov 29, 2022
2.4k Dec 31, 2022
568 Dec 14, 2022