Skip to content

Latest commit

 

History

History
 
 

toolkit

Folders and files

NameName
Last commit message
Last commit date

parent directory

..
data
data
 
 
lit_models
lit_models
 
 
models
models
 
 
scripts
scripts
 
 
.gitignore
.gitignore
 
 
README.md
README.md
 
 
eval.py
eval.py
 
 
main.py
main.py
 
 
requirements.txt
requirements.txt
 
 

README.md

A Prompt Learning Framework for Knowledge Graph with Pre-trained Language Models.

Quick Start

Installation

Step1

Create the conda env.

conda create -n promptkg python=3.8

Step2

Install the dependence.

pip install -r requirements.txt

Step3

Install our preprocessed datasets and put them into the dataset folder.

Dataset (KGC) Google Drive Baidu Cloud
WN18RR google drive baidu drive axo7
FB15k-237 google drive baidu drive ju9t
MetaQA google drive baidu drive hzc9
KG20C google drive baidu drive stnh

Run your first experiment

We provide four tasks in our toolkit as Knowledgeg Graph Completion (KGC), Question Answering (QA), Recomandation (REC) and LAnguage Model Analysis (LAMA).

  • KGC is our basic task to the knowledge graph embedding and evaluate the ability of the models. ** You can run the script under kgc folder to train the model and get the KG embeddings (take simkgc as example).

    bash ./scripts/kgc/simkgc.sh

  • For QA task, you can run the script files under metaqa. ** We suggest you use generative model to solve the QA task as below:

    bash ./scripts/metaqa/run.sh

  • For REC task, you need to firstly get the KG embeddings and then train the rec system models. ** use two-stage scripts below:

    bash ./scripts/kgrec/pretrain_item.sh
bash ./scripts/kgrec/ml20m.sh

  • For LAMA task, you can use the files under lama. ** We provide BERT and RoBERTa PLMs to evaluate their performance and with our KG embeddings (plet).

    bash ./scripts/lama/lama_roberta.sh

Implemented Models

Models Knowledge Graph Completion Question Answering Recomandation LAnguage Model Analysis
KG-BERT ✔ ✔
GenKGC ✔
KGT5 ✔ ✔
kNN-KGE ✔ ✔ ✔
SimKGC ✔

Framework

Process on your own data

For each knowledge graph, we have 5 files.

  • train.tsv, dev.tsv, test.tsv, list as (h, r, t) for entity id and relation id (start from 0).
  • entity2text.txt, as (entity_id, entity description).
  • relation2text.txt , as (relation_id, relation description).

For downstream tasks

TODO

  • add essemble-model for using
  • add more kgc models based on pretrained language models
  • add soft prompt in base model.
  • add more dataset with the same format.