jeongrok/VGCN-BERT

VGCN-BERT: Augmenting BERT with Graph Embedding for Text Classification

Authors: Zhibin Lu (zhibin.lu@umontreal.ca), Pan Du (pan.du@umontreal.ca), Jian-Yun Nie (nie@iro.umontreal.ca)

Overview

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},
   }

Requirements

• Python >= 3.7
• NLTK 3.4
• PyTorch >= 1.0
• Huggingface transformer 0.6.2
• SST-2
• CoLA

Directory Structure

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.

Pretrained BERT weights

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

Running the code

Pre-processing datasets

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

Train models

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