SRTK is a toolkit for semantic-relevant subgraph retrieval from large-scale knowledge graphs. It currently supports Wikidata, Freebase and DBpedia.
A minimum walkthrough of the retrieve process:
pip install srtk
There are mainly five subcommands of SRTK, which covers the whole pipeline of subgraph retrieval.
For retrieval:
srtk link
: Link entity mentions in texts to a knowledge graph. Currently Wikidata and DBpedia are supported out of the box.srtk retrieve
: Retrieve semantic-relevant subgraphs from a knowledge graph with a trained retriever. It can also be used to evaluate a trained retriever.srtk visualize
: Visualize retrieved subgraphs using a graph visualization tool.
For training a retriever:
srtk preprocess
: Preprocess a dataset for training a subgraph retrieval model.srtk train
: Train a subgraph retrieval model on a preprocessed dataset.
Use srtk [subcommand] --help
to see the detailed usage of each subcommand.
srtk retrieve [-h] -i INPUT -o OUTPUT [-e SPARQL_ENDPOINT] -kg {freebase,wikidata,dbpedia}
-m SCORER_MODEL_PATH [--beam-width BEAM_WIDTH] [--max-depth MAX_DEPTH]
[--evaluate] [--include-qualifiers]
The scorer-model-path
argument can be any huggingface pretrained encoder model. If it is a local
path, please ensure the tokenizer is also saved along with the model.
srtk visualize [-h] -i INPUT -o OUTPUT_DIR [-e SPARQL_ENDPOINT]
[-kg {wikidata,freebase}] [--max-output MAX_OUTPUT]
A scorer is the model used to navigate the expanding path. At each expanding step, relations scored higher with scorer are picked as relations for the next hop.
The score is based on the embedding similarity of the to-be-expanded relation with the query (question + previous expanding path).
The model is trained in a distant supervised learning fashion. Given the question entities and the answer entities, the model uses the shortest paths along them as the supervision signal.
-
prepare training samples where question entities and answer entities are know.
The training data should be saved in a jsonl file (e.g.
data/grounded.jsonl
). Each training sample should come with the following format:{ "id": "sample-id", "question": "Which universities did Barack Obama graduate from?", "question_entities": [ "Q76" ], "answer_entities": [ "Q49122", "Q1346110", "Q4569677" ] }
-
Preprocess the samples with
srtk preprocess
command.srtk preprocess [-h] -i INPUT -o OUTPUT [--intermediate-dir INTERMEDIATE_DIR] -e SPARQL_ENDPOINT -kg {wikidata,freebase} [--search-path] [--metric {jaccard,recall}] [--num-negative NUM_NEGATIVE] [--positive-threshold POSITIVE_THRESHOLD]
Under the hood, it does four things:
- Find the shortest paths between the question entities and the answer entities.
- Score the searched paths with Jaccard scores with the answers.
- Negative sampling. At each expanding step, the negative samples are those false relations connected to the tracked entities.
- Generate training dataset as a jsonl file.
The scorer should be initialized from a pretrained encoder model from huggingface hub. Here I used intfloat/e5-small
, which is a checkpoint of the BERT model.
srtk train --data-file data/train.jsonl \
--model-name-or-path intfloat/e5-small \
--save-model-path artifacts/scorer
SRTK is compatible with any language encoder or encoder-decoder models from huggingface hub. You only need to specify the model name or path for arguments like --model-name-or-path
or --scorer-model-path
.
Here we provide some trained models for subgraph retrieval.
Model | Dataset | Base Model | Notes |
---|---|---|---|
drt/srtk-scorer |
WebQSP, SimpleQuestionsWikidata, SimpleDBpediaQA | roberta-base |
Jointly trained for Wikidata, Freebase and DBpedia. |
- End-to-end Subgraph Retrieval
- Train a Retriever on Wikidata with Weak Supervision
- Train a Retriever on Freebase with Weak Supervision
- Supervised Training with Wikidata Simple Questions
- Extend SRTK to other Knowledge Graphs
@article{shen2023srtk,
title={SRTK: A Toolkit for Semantic-relevant Subgraph Retrieval},
author={Shen, Yuanchun},
journal={arXiv preprint arXiv:2305.04101},
year={2023}
}
@inproceedings{zhang2022subgraph,
title={Subgraph Retrieval Enhanced Model for Multi-hop Knowledge Base Question Answering},
author={Zhang, Jing and Zhang, Xiaokang and Yu, Jifan and Tang, Jian and Tang, Jie and Li, Cuiping and Chen, Hong},
booktitle={Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
pages={5773--5784},
year={2022}
}
This project is licensed under the terms of the MIT license.