TimeR4

Abstract Temporal Knowledge Graph Question Answering (TKGQA) aims to answer temporal questions using knowledge in Temporal Knowledge Graphs (TKGs). Previous works employ pre-trained TKG embeddings or graph neural networks to incorporate the knowledge of TKGs. However, these methods fail to fully understand the complex semantic information of time constraints. In contrast, Large Language Models (LLMs) have shown exceptional performance in knowledge graph reasoning, unifying both semantic understanding and structural reasoning. To further enhance LLMs’ temporal reasoning ability, this paper aims to integrate temporal knowledge from TKGs into LLMs through a Time-aware Retrieve-Rewrite-Retrieve-Rerank framework, which we named TimeR4. Specifically, to reduce temporal hallucination in LLMs, we propose a retrieve-rewrite module to rewrite questions using background knowledge stored in the TKGs, thereby acquiring explicit time constraints. Then, we implement a retrieve-rerank module aimed at retrieving semantically and temporally relevant facts from the TKGs and reranking according to the temporal constraints. To achieve this, we fine-tune a retriever using the contrastive time-aware learning framework. Our approach achieves great improvements, with relative gains of 47.8% and 22.5% on two datasets, underscoring its effectiveness in boosting the temporal reasoning abilities of LLMs.

Codes for "TimeR4: Time-aware Retrieval-Augmented Large Language Models for Temporal Knowledge Graph Question Answering".

We have also made further extensions and improvements to this work, which are presented in "Plan of Knowledge: Retrieval-Augmented Large Language Models for Temporal Knowledge Graph Question Answering".

Requirements

pip install -r requirements.txt

datasets

unzip datasets.zip
unzip results.zip

Pre-trained weights

You can download the pre-trained weights of retriever and LLMs here..

Inference

We run our code on 1 NVIDIA A6000 GPUs for inference and 2 NVIDIA A6000 GPUs for training.

Step1: Retrieve-rewrite-retrieve-rerank

python main.py --output_path  datasets/MultiTQ/results/test_prompt.json --question_path datasets/MultiTQ/questions/test.json --retrieve_name models/multi_model --triplet_path datasets/MultiTQ/kg/full.txt

Step2: Reasoning

python predict_answer.py  --model_path  model/LLMs/Llama2_MultiTQ/checkpoint-1800  -d datasets/MultiTQ/prompt/test_prompt.json --debug --predict_path datasets/MultiTQ/result

Training

Training time-aware retrieval

1. Constructing negative samples.

python construct_negatives.py --output_path  datasets/MultiTQ/negatives.json  --question_path datasets/MultiTQ/questions/train.json --datasets multi --entity2id_path datasets/MultiTQ/kg/entity2id.json --time2id_path datasets/MultiTQ/kg/ts2id.json --triplet_path datasets/MultiTQ/kg/full.txt --relation2id_path datasets/MultiTQ/kg/relation2id.json

2. Fine-tune the LMs.

python fine_tuned_retriever.py --path datasets/multi/negatives.json --model_name time_aware_model_tp

Fine-tuning Llama2

We use Llama-X to fine tune llama2.

deepspeed --master_port 25678 --include localhost:0,1 train.py \
    --model_name_or_path llama-7b-chat-hf  \
    --data_path dataset/MultiTQ/prompt/train_prompt.json \
    --output_dir Llama2_MultiTQ  \
    --num_train_epochs 3 \
    --model_max_length 1024 \
    --per_device_train_batch_size 256 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 40 \
    --save_total_limit 1 \
    --learning_rate 2e-5 \
    --warmup_steps 2 \
    --logging_steps 2 \
    --lr_scheduler_type "cosine" \
    --report_to "tensorboard" \
    --gradient_checkpointing True \
    --deepspeed deepspeed_config.json \
    --bf16 True 

Results

We provide the results of ChatGPT, Llama2, and fine-tuned Llama2 for both the input question and the prompt generated by the $TimeR^4$ framework in the results folder. Fine-tuned Llama2 is the results we report in the paper.

Bibinfo

If you found this repo helpful, please help us by citing this paper:

@inproceedings{qian-etal-2024-timer4,
    title = "{T}ime{R}$^4$ : Time-aware Retrieval-Augmented Large Language Models for Temporal Knowledge Graph Question Answering",
    author = "Qian, Xinying  and
      Zhang, Ying  and
      Zhao, Yu  and
      Zhou, Baohang  and
      Sui, Xuhui  and
      Zhang, Li  and
      Song, Kehui",
    editor = "Al-Onaizan, Yaser  and
      Bansal, Mohit  and
      Chen, Yun-Nung",
    booktitle = "Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2024",
    address = "Miami, Florida, USA",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2024.emnlp-main.394",
    doi = "10.18653/v1/2024.emnlp-main.394",
    pages = "6942--6952",
}