CoKE

CoKE: Contextualized Knowledge Graph Embedding

Introduction

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.

Requirements

The code has been tested running under the following environments:

• Python 3.6.5 with the following dependencies:

  • PaddlePaddle 1.6.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.

Model Training and Evaluation

step1. Download dataset files

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

step2. Data preprocess

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

step3. Training

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

step4. Evaluation

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

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

Reproducing the results

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

Citation

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 and License

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.