NuTrea

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.

Setup

1. 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

2. 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/

3. (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.

Run

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.

UPDATES

2023.12 Initial code release

Citation

@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}
}

License

Code is released under MIT License.