Official implementation of NeurIPS 2023 paper, "NuTrea: Neural Tree Search for Context-guided Multi-hop KGQA" by Hyeong Kyu Choi, Seunghun Lee, Jaewon Chu, and Hyunwoo J. Kim.
- Environment
git clone https://github.com/mlvlab/NuTrea.git
cd NuTrea
conda create -n nutrea python=3.8
conda activate nutrea
conda install pytorch==1.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia
conda install pyg -c pyg
pip install -r src/requirements.txt
- Datasets
Download dataset here. Unzip data files and place under 'data/' directory as
NuTrea/
├── src/
├── checkpoints/
├── ...
└── data/
├── webqsp/
├── CWQ/
├── metaqa-1hop/
├── metaqa-2hop/
└── metaqa-3hop/
- (optional) Checkpoints
Download the pretrained checkpoint(s) and place under "checkpoints/" directory.
| Dataset | Metric | Hit@1 | F1 | Checkpoint |
|---|---|---|---|---|
| WebQuestionsSP | Hit@1 | 77.43 | 71.01 | drive |
| WebQuestionsSP | F1 | 76.88 | 72.70 | drive |
| ComplexWebQuestions | Hit@1 | 53.61 | 49.41 | drive |
| ComplexWebQuestions | F1 | 53.16 | 49.53 | drive |
| MetaQA-1hop | Hit@1 | 97.40 | 97.53 | drive |
| MetaQA-1hop | F1 | 97.25 | 97.62 | drive |
| MetaQA-2hop | Hit@1 | 99.99 | 99.82 | drive |
| MetaQA-2hop | F1 | 99.99 | 99.82 | drive |
| MetaQA-3hop | Hit@1 | 98.89 | 87.06 | drive |
| MetaQA-3hop | F1 | 98.89 | 87.06 | drive |
- Note, MetaQA-2hop and -3hop Hit@1 and F1 model checkpoints contain identical parameters, respectively.
- Also note that train / eval performance may vary across different GPU and environments.
Runnable scripts are in the "runs/" directory. To run evaluation with checkpoints,
sh runs/wqp_eval.sh
To train the model,
sh runs/wqp_train.sh
Run different scripts to test different datasets and parameters.
2023.12 Initial code release
@inproceedings{choi2023nutrea,
title={NuTrea: Neural Tree Search for Context-guided Multi-hop KGQA},
author={Choi, Hyeong Kyu and Lee, Seunghun and Chu, Jaewon and Kim, Hyunwoo J.},
booktitle={Advances in Neural Information Processing Systems},
year={2023}
}
Code is released under MIT License.