TRIO: Task Refinement with Individual Optimization for Retrieval-Augmented Generation

🌟 Introduction

This repository is the official implementation of TRIO (Task Refinement with Individual Optimization for Retrieval-Augmented Generation), a novel framework designed to improve Retrieval-Augmented Generation (RAG) systems.

---

🔧 Installation

Clone this repository, then create a conda environment and install the required packages.

# Clone repository
git clone https://github.com/yingtaoren/TRIO.git
cd TRIO

# Create conda environment
conda create -n trio python=3.10
conda activate trio

# Install packages
pip install -r requirements.txt

💡 Note: If you encounter any issues when installing the Python packages, we recommend following the official installation instructions provided by FlashRAG#Installation.

Additional Dependencies

For faiss installation, use conda:

# CPU-only version
conda install -c pytorch faiss-cpu=1.8.0

# GPU(+CPU) version
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

---

📊 Dataset

The datasets used in this project follow the same format as those pre-processed by FlashRAG. All datasets are available at 🤗 Huggingface datasets.

Supported Datasets

Type	Dataset
Multi-hop QA	HotpotQA
Multi-hop QA	2WikiMultiHopQA
Multi-hop QA	MuSiQue
Multi-hop QA	Bamboogle

Dataset Setup

After downloading the datasets from FlashRAG, create a /dataset folder in the project directory and place the downloaded data inside:

TRIO
├── config
├── dataset
│   ├── 2wiki
│   ├── bamboogle
│   ├── HotpotQA
│   ├── musique
├── model
├── scripts
├── main.py
├── README.md
└── requirements.txt

💡 Note: If you wish to use a custom dataset path, modify the data_dir field in config/base_config.yaml.

---

📚 Document Corpus & Index

We use the wiki18_100w dataset as the document corpus. Both the document corpus and the index can be downloaded from:

• ModelScope Dataset (look in retrieval_corpus folder)

💡 Note: If you wish to use a custom corpus path, modify the index_path and corpus_path fields in config/base_config.yaml.

---

🤖 Models

Retriever Model

We use e5-base-v2 as the default retriever. Download from: 🤗 intfloat/e5-base-v2

Generator Models

This project supports all LLMs compatible with HuggingFace and vLLM. Recommended models:

Model	Link
Llama-3.1-8B-Instruct	🔗 Link
Qwen3-4B-Instruct	🔗 Link

We also provided our trained moedls here.

Please specify the path to your downloaded model using the model2path field in config/base_config.yaml:

model2path:
  e5: "/path/to/e5-base-v2"
  llama3-8B-instruct: "/path/to/Llama-3.1-8B-Instruct"
  # Add your custom models here

---

🏋️ Training

Training Data Collection

For training data collection, please refer to the DRAG repository which provides the data collection pipeline.

Supervised Fine-tuning (SFT)

For the judge agent training, run the SFT training script:

bash sft_train.sh

Reinforcement Learning (RL)

For query-extension and answer agent training, we use OpenRLHF. Run the training script:

bash my_train.sh

📈 Evaluation

Running Evaluation

Evaluate TRIO on the supported datasets:

python main.py --method_name "TRIO" \
               --gpu_id "0" \
               --dataset_name "HotpotQA" \
               --generator_model "llama3-8B-instruct-TRIO" \
               --max_query_debate_rounds 3

---

✨ Acknowledgments

We gratefully acknowledge the following projects:

• FlashRAG: A Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research.
• DRAG: Debate-Augmented RAG framework that inspired our work.
• OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework.

---