Code for our paper "Reasoning Through Memorization: Nearest Neighbor Knowledge Graph Embeddings".
To install requirements:
pip install -r requirements.txt
Use the command below to add entities to BERT and train the entity embedding layer to use in the later training. For another dataset WN18RR
just replacing the dataset name will be fine.
./scripts/pretrain_fb15k.sh
The parameters of Entity Embedding Layer trained will be used in the next Entity prediction task
.
Use the command below to train the model to predict the correct entity in the masked position.
./scripts/fb15k-237/fb15k.sh
After training the model in Entity prediction task
, we use the model to get the knowledge store built from triples and descriptions.
./scripts/fb15k-237/get_knowledge_store.sh
Here we have a trained model and our knowledge store, use the command below to inference in the test set.
./scripts/fb15k-237/inference.sh
And for inductive setting, the command is similar to the transductive setting (just replace the dataset
with inductive dataset), the code will automatically handle the differences.