Introduction : This repo is for the final project of NYU DS-1005 Inference and Representations for Bofei Zhang, Yunan Hu, Kaitai Zhang. We adopt Knowledge Base Attention Network (KBAT) on a novel knowledge graph completion task, Drug-Drug Interaction task.
The folder inferences
contains code that can processed drugbank data and prepared the training, testing data for KBAT training and evaluation. You can run this based on the following order:
# make sure you have the original dataset from here https://www.drugbank.ca/releases/latest
# if run on the prince, please request more cpu resources, otherwise it has a memory issue
python process.py
# this script filter down the original dataset into a smaller one
python filter_down.py
# if builds a folder that contains all data for training
python create_dataset_kbat.py
Move the generated folder to data
, then you are ready to go.
To reproduce the experiment, see the scripts in scripts/
folder. We ran all experiments on NYU HPC Prince. Below is an example to run an experiment, see main.py
for details about arguments:
python main.py python main.py --data ./data/drugbank_1861/ --pretrained_emb False\
--output_folder ./checkpoints/drugbank1861_tanh/ \
--epochs_gat 1200\
--epochs_conv 200\
--batch_size_conv 512 --batch_size_gat 80000\
--tanh yes
Below are original README file.
Source code for our ACL 2019 paper: Learning Attention-based Embeddings for Relation Prediction in Knowledge Graphs Blog link for this publication.
Please download miniconda from above link and create an environment using the following command:
conda env create -f pytorch35.yml
Activate the environment before executing the program as follows:
source activate pytorch35
We used five different datasets for evaluating our model. All the datasets and their folder names are given below.
- Freebase: FB15k-237
- Wordnet: WN18RR
- Nell: NELL-995
- Kinship: kinship
- UMLS: umls
Parameters:
--data
: Specify the folder name of the dataset.
--epochs_gat
: Number of epochs for gat training.
--epochs_conv
: Number of epochs for convolution training.
--lr
: Initial learning rate.
--weight_decay_gat
: L2 reglarization for gat.
--weight_decay_conv
: L2 reglarization for conv.
--get_2hop
: Get a pickle object of 2 hop neighbors.
--use_2hop
: Use 2 hop neighbors for training.
--partial_2hop
: Use only 1 2-hop neighbor per node for training.
--output_folder
: Path of output folder for saving models.
--batch_size_gat
: Batch size for gat model.
--valid_invalid_ratio_gat
: Ratio of valid to invalid triples for GAT training.
--drop_gat
: Dropout probability for attention layer.
--alpha
: LeakyRelu alphas for attention layer.
--nhead_GAT
: Number of heads for multihead attention.
--margin
: Margin used in hinge loss.
--batch_size_conv
: Batch size for convolution model.
--alpha_conv
: LeakyRelu alphas for conv layer.
--valid_invalid_ratio_conv
: Ratio of valid to invalid triples for conv training.
--out_channels
: Number of output channels in conv layer.
--drop_conv
: Dropout probability for conv layer.
To reproduce the results published in the paper:
When running for first time, run preparation script with:
$ sh prepare.sh
-
Wordnet
$ python3 main.py --get_2hop True
-
Freebase
$ python3 main.py --data ./data/FB15k-237/ --epochs_gat 3000 --epochs_conv 150 --weight_decay_gat 0.00001 --get_2hop True --partial_2hop True --batch_size_gat 272115 --margin 1 --out_channels 50 --drop_conv 0.3 --output_folder ./checkpoints/fb/out/
Please cite the following paper if you use this code in your work.
@InProceedings{KBGAT2019,
author = "Nathani, Deepak and Chauhan, Jatin and Sharma, Charu and Kaul, Manohar",
title = "Learning Attention-based Embeddings for Relation Prediction in Knowledge Graphs",
booktitle = "Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics",
year = "2019",
publisher = "Association for Computational Linguistics",
location = "Florence, Italy",
}
For any clarification, comments, or suggestions please create an issue or contact [email protected]