Code for the NLPCC2023 paper "Reasoning Through Memorization: Nearest Neighbor Knowledge Graph Embeddings with Language Models".

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 (e.g., faiss.dump file), 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.