This is the PaddlePaddle implementation of the CoKE model for Knowledge Graph Embedding(KGE).
CoKE is a novel KGE paradigm that learns dynamic, flexible, and fully contextualized entity and relation representations for a given Knowledge Graph(KG). It takes a sequence of entities and relations as input, and uses Transformer to obtain contextualized representations for its components. These representations are hence dynamically adaptive to the input, capturing contextual meanings of entities and relations therein.
Evaluation on a wide variety of public benchmarks verifies the superiority of CoKE in link prediction (also known as Knowledge Graph Completion, or KBC for short) and path query answering tasks. CoKE performs consistently better than, or at least equally well as current state-of-the-art in almost every case.
The code has been tested running under the following environments:
-
Python 3.6.5 with the following dependencies:
- PaddlePaddle 1.5.0
- numpy 1.16.3
-
Python 2.7.14 for data_preprocess
-
GPU environments:
- CUDA 9.0, CuDNN v7 and NCCL 2.3.7
- GPU: all the datasets run on 1 P40 GPU with our given configurations.
Download dataset files used in our paper by running:
sh wget_datasets.sh
This will first download the 4 widely used KBC datasets (FB15k&WN18, FB15k-237, WN18RR) and 2 path query answering datasets (wordnet_paths and freebase_paths) .
Then it organize the train/valid/test files as the following data
directory:
data
├── fb15k
│ ├── test.txt
│ ├── train.txt
│ └── valid.txt
├── fb15k237
│ ├── test.txt
│ ├── train.txt
│ └── valid.txt
├── pathqueryFB #the original data name is: freebase_paths
│ ├── dev
│ ├── test
│ └── train
├── pathqueryWN #the original data name is: wordnet_paths
│ ├── dev
│ ├── test
│ └── train
├── wn18
│ ├── test.txt
│ ├── train.txt
│ └── valid.txt
└── wn18rr
├── test.txt
├── train.txt
└── valid.txt
Data preprocess commands are given in data_preprocess.sh
.
It takes raw train/valid/test files as input, and generates CoKE training and evaluation files.
sh data_preprocess.sh
Model training commands are given in kbc_train.sh
for KBC datasets, and pathquery_train.sh
for pathquery datasets.
These scripts take a configuration file and GPU-ids as input arguments.
Train the model with a given configuration file.
For example, the following commands train fb15k and pathqueryFB each with a configuration file:
sh kbc_train.sh ./configs/fb15k_job_config.sh 0
sh pathquery_train.sh ./configs/pathqueryFB_job_config.sh 0
Model evaluation commands are given in kbc_test.sh
for KBC datasets, and pathquery_test.sh
for pathquery datasets.
These scripts take a configuration file and GPU-ids as input arguments.
For example, the following commands evaluate on fb15k and pathqueryFB:
sh kbc_test.sh ./configs/fb15k_job_config.sh 0
sh pathquery_test.sh ./configs/pathqueryFB_job_config.sh 0
We also provide trained model checkpoints on the 4 KBC datasets. Download these models to kbc_models
directory using the following command:
sh wget_kbc_models.sh
The kbc_models
contains the following files:
kbc_models
├── fb15k
│ ├── models
│ └── vocab.txt #md5: 0720db5edbda69e00c05441a615db152
├── fb15k237
│ ├── models
│ └── vocab.txt #md5: e843936790e48b3cbb35aa387d0d0fe5
├── wn18
│ ├── models
│ └── vocab.txt #md5: 4904a9300fc3e54aea026ecba7d2c78e
└── wn18rr
├── models
└── vocab.txt #md5: c76aecebf5fc682f0e7922aeba380dd6
Check that your preprocessed vocab.txt
files are identical to ours before evaluation with these models.
Results on KBC datasets:
Dataset | MRR | HITS@1 | HITS@3 | HITS@10 |
---|---|---|---|---|
FB15K | 0.855 | 0.826 | 0.872 | 0.906 |
FB15K237 | 0.364 | 0.272 | 0.400 | 0.549 |
WN18 | 0.952 | 0.947 | 0.955 | 0.960 |
WN18RR | 0.484 | 0.450 | 0.496 | 0.553 |
Results on path query datasets:
Dataset | MQ | HITS@10 |
---|---|---|
Freebase | 0.950 | 0.777 |
WordNet | 0.942 | 0.679 |
Here are the configs to reproduce our results.
These are also given in the configs/${TASK}_job_config.sh
files.
Dataset | NetConfig | epoch | dropout | label smoothing |
---|---|---|---|---|
FB15K | L=12, H=256, A=4 | 400 | 0.2 | 0.8 |
WN18 | L=12, H=256, A=4 | 600 | 0.2 | 0.2 |
FB15K237 | L=12, H=256, A=4 | 500 | 0.5 | 0.25 |
WN18RR | L=12, H=256, A=4 | 1000 | 0.3 | 0.15 |
pathqueryFB | L=12, H=256, A=4 | 10 | 0.1 | 1 |
pathqueryWN | L=12, H=256, A=4 | 5 | 0.1 | 1 |
If you use any source code included in this project in your work, please cite the following paper:
@article{wang2019:coke,
title={CoKE: Contextualized Knowledge Graph Embedding},
author={Wang, Quan and Huang, Pingping and Wang, Haifeng and Dai, Songtai and Jiang, Wenbin and Liu, Jing and Lyu, Yajuan and Wu, Hua},
journal={arXiv:1911.02168},
year={2019}
}
Copyright 2019 Baidu.com, Inc. All Rights Reserved Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.