Authors: Zhibin Lu (zhibin.lu@umontreal.ca), Pan Du (pan.du@umontreal.ca), Jian-Yun Nie (nie@iro.umontreal.ca)
This is the implementation of our ECIR 2020 paper: "VGCN-BERT: Augmenting BERT with Graph Embedding for Text Classification".
If you make use of this code or the VGCN-BERT approach in your work, please cite the following paper:
@inproceedings{ZhibinluGraphEmbedding,
author = {Zhibin Lu and Pan Du and Jian-Yun Nie},
title = {VGCN-BERT: Augmenting BERT with Graph Embedding for Text Classification},
booktitle = {Advances in Information Retrieval - 42nd European Conference on {IR}
Research, {ECIR} 2020, Lisbon, Portugal, April 14-17, 2020, Proceedings,
Part {I}},
series = {Lecture Notes in Computer Science},
volume = {12035},
pages = {369--382},
publisher = {Springer},
year = {2020},
}
- Python >= 3.7
- NLTK 3.4
- PyTorch >= 1.0
- Huggingface transformer 0.6.2
- SST-2
- CoLA
The data/ directory contains the original datasets and the corresponding pre-processed data. The output/ directory saves the model files while training the models. The pytorch_pretrained_bert/ directory is the module of Huggingface transformer 0.6.2. The other files with the .py suffix are models, utility functions, training programs, and prediction programs.
We mainly use https://huggingface.co/bert-base-uncased as the pretrained Bert weights. If you want first download the model fiels or use yourself pretrained model, you can modify the 2 env variables in .env
as follows.
TRANSFORMERS_OFFLINE=1
HUGGING_LOCAL_MODEL_FILES_PATH=/tmp/local-huggingface-models
Run prepare_data.py to pre-process the dataset and generate the vocabulary graph.
Examples:
cd VGCN-BERT
python -m vgcn_bert.prepare_data --ds cola
Run train_vanilla_vgcn_bert.py to train the Vanilla-VGCN-BERT model. Before training the model, you must ensure that you have pre-processed the dataset.
Examples:
cd VGCN-BERT
python -m vgcn_bert.train_vanilla_vgcn_bert --ds cola
Run train_vgcn_bert.py to train the VGCN-BERT model. Before training the model, you must ensure that you have pre-processed the dataset.
Examples:
cd VGCN-BERT
python -m vgcn_bert.train_vgcn_bert --ds cola