From Retrieval to Translation: Translating Query into Graph-level Clues for Retrieval-Augmented Generation

📖 Overview

⚙️ Environment

1. Install requirements:

conda create -n KG-Translator python=3.9
conda activate KG-Translator
pip install -r requirements.txt

2. Install SEAL for constrained indexing:

git clone --recursive https://github.com/facebookresearch/SEAL.git
cd SEAL
env CFLAGS='-fPIC' CXXFLAGS='-fPIC' res/external/sdsl-lite/install.sh
pip install -r requirements.txt
pip install -r requirements_extra.txt
pip install -e .

SEAL needs some patches for newer transformers versions. Edit seal/beam_search.py:

• Change from transformers.generation_utils import xxx to from transformers.generation.utils import xxx
• Change from transformers.generation_logits_process import TopKLogitsWarper to from transformers.utils.dummy_pt_objects import TopKLogitsWarper

3. Structure-aware Fine-tuning environment:

Structure-aware Fine-tuning uses the official environment of LLaMA-Factory. For environment installation, please refer to LLaMA-Factory.

📁 Data

Data for Parsing-based KG Construction, KG-constrained Indexing and Structure-aware Fine-tuning should be put into data folder. We provide an example for each stage.

• data/collection is the corpus.
• data/ner has the results of NER for ParseKG construction and KG-constrained index.
• data/ground_truth is the ground truth for retrieval and QA.
• data/test is the testing data.
• data/train is the training data for fine-tuning and is in the format required by LLaMA-Factory.

🔍 KG-Translator

Parsing-based KG Construction

Run corpus_parsing.py to construct ParseKG:

python corpus_parsing.py \
  --dataset_name "2wikimqa" \
  --dataset_type "2h_bridge" \
  --worker_nums 16

Structure-constrained KG Translation

1. KG-constrained Indexing. Run kg_constrained_indexing.py to construct the constrained index:

python kg_constrained_indexing.py \
  --dataset_name "2wikimqa" \
  --dataset_type "2h_bridge" \
  --model_name_or_path "YOUR_TOKENIZER_PATH" \
  --emb_model_dir "YOUR_EMBEDDING_MODEL_DIR" \
  --emb_model_name "all-MiniLM-L6-v2"

2. Structure-aware Fine-tuning.

Copy data/train to LLaMA-Factory/data. Paste the configuration below into data/dataset_info.json to register the training data:

"kg_translator": {
    "file_name": "train/train_llama_factory_format.jsonl"
  },

Then run the script below to fine-tune KG-Translator:

conda activate llama_factory

export NCCL_P2P_DISABLE=1

KG_TRANSLATOR_OUTPUT_PATH = "YOUR_KG_TRANSLATOR_PATH"

CUDA_VISIBLE_DEVICES=0 llamafactory-cli train \
    --model_name_or_path meta-llama/Llama-3.2-1B-Instruct \
    --trust_remote_code \
    --stage sft \
    --do_train \
    --finetuning_type full \
    --dataset kg_translator  \
    --template llama3 \
    --cutoff_len 2048 \
    --max_samples 50000 \
    --overwrite_cache \
    --preprocessing_num_workers 16 \
    --dataloader_num_workers 4 \
    --output_dir KG_TRANSLATOR_OUTPUT_PATH \
    --logging_steps 3 \
    --save_steps 500 \
    --plot_loss \
    --overwrite_output_dir \
    --save_only_model true \
    --report_to none \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 4 \
    --learning_rate 1e-5 \
    --num_train_epochs 1.0 \
    --lr_scheduler_type cosine \
    --warmup_ratio 0.1 \
    --seed 42 \
    --bf16 \
    --ddp_timeout 180000000 \
    --deepspeed examples/deepspeed/ds_z2_config.json

3. Obtain the KG-Translator results. Run rag_pipeline.py to obtain the retrieval and QA results of KG-Translator:

python rag_pipeline.py \
  --dataset_name "2wikimqa" \
  --dataset_type "2h_bridge" \
  --test_num 1000 \
  --emb_model_name "all-MiniLM-L6-v2" \
  --retriever_path "YOUR_KG_TRANSLATOR_PATH" \
  --generator_name "YOUR_GENERATOR_NAME" \
  --note "YOUR_EXPERIMENT_NOTE" \
  --head_num 3 \
  --relation_num 1 \
  --tail_num 0

📊 Evaluation

Run evaluate.py to evalute the final results.

python evaluate.py --log_dir "./output"