Home

Awesome

KG-R3

Code for the CIKM'23 paper "A Retrieve-and-Read Framework for Knowledge Graph Link Prediction" ACM link

<!-- Code will be released soon. -->

A Retrieve-and-Read Framework for Knowledge Graph Link Prediction (KG-R3)

KG Link Prediction Results

DatasetMRRHITS@1HITS@3HITS@10
FB15K-237.390.315.413.539
WN18RR.472.439.481.537

Install dependencies

  1. Create a new conda virtual env

  2. Install horovod

HOROVOD_WITH_PYTORCH=1 --no-cache-dir --ignore-installed pip install horovod[pytorch] --extra-index-url https://download.pytorch.org/whl/cu113
  1. Install other dependencies
pip install -r requirements.txt

Download data

Download the preprocessed subgraphs and KG triples from this link from respective directories FB15K-237 and WN18RR are place them in a data/ directory.

Dump retriever subgraphs (optional)

Preprocess data

pickle dataloader batches for faster training

FB15K-237, Minerva retriever

python -u dump_preproc_data.py --dataset-path data/FB15K-237/ \
--sampling-type minerva \
--batch-size 256 --out-dir data/FB15K-237/train_preproc/ \
--graph-connection type_1 --split train

Training

FB15K-237

train, Minerva retriever

python -u main.py --dataset-path data/FB15K-237/ --cuda \
--save-dir ckpts/CKPT_DIR/ --sampling-type minerva \
--lr 1e-2 --warmup 0.1 --batch-size 512 \
--n-epochs 300 --patience 20 \
--seed 12548 > ckpts/CKPT_DIR/log.txt 2>&1

WN18RR

train, Minerva retriever

python -u main.py --dataset-path data/WN18RR/ --cuda \
--save-dir ckpts/CKPT_DIR/ --sampling-type minerva \
--lr 0.00175 --label-smoothing 0.1 --warmup 0.1 \
--batch-size 256 --n-epochs 500 \
--patience 100 --beam-size 40 --add-segment-embed --add-inverse-rels \
--seed 12548 > ckpts/CKPT_DIR/log.txt 2>&1

Evaluation (specify split)

python eval.py --dataset-path <DATA_PATH> --cuda \
--ckpt-path ckpts/CKPT_DIR/model.pt \
--split <valid/test> --sampling-type minerva \
--graph-connection type_1  \
[--beam-size <>] [--add-segment-embed] [--add-inverse-rels]

Citation

@inproceedings{10.1145/3583780.3614769,
author = {Pahuja, Vardaan and Wang, Boshi and Latapie, Hugo and Srinivasa, Jayanth and Su, Yu},
title = {A Retrieve-and-Read Framework for Knowledge Graph Link Prediction},
year = {2023},
isbn = {9798400701245},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3583780.3614769},
doi = {10.1145/3583780.3614769},
pages = {1992–2002},
numpages = {11},
location = {Birmingham, United Kingdom},
series = {CIKM '23}
}