WisdomShell/GRIP-Llama-3-8B

Retrieval as Generation: A Unified Framework with Self-Triggered Information Planning

English | 简体中文

[ACL'26 Main Conference]

Bo Li  Mingda Wang  GeXiang Fang  Shikun Zhang  Wei Ye 

Traditional RAG (Retrieval-Augmented Generation) systems treat retrieval as an external, one-shot intervention, rigidly fetching documents before generation begins, which often fails when information needs emerge gradually during complex reasoning. Even dynamic search methods heavily rely on disconnected external controllers or heuristic rules. We believe that, much like human cognitive processes, retrieval should be an intrinsic, generative capability. LLMs must be able to autonomously evaluate their knowledge, trigger searches, and formulate contextual follow-up queries tightly coupled with their evolving reasoning states. GRIP (Generation-guided Retrieval with Information Planning) embodies this new paradigm. Under the framework of Retrieval as Generation, our model internalizes retrieval decisions directly into token-level decoding using specific control tokens. This approach shifts from relying on auxiliary multi-stage search modules to achieving end-to-end, self-triggered information planning within a single autoregressive trajectory.

🌟 Key Features

• 🎯 Token-Driven Control: Embeds retrieval behaviors directly into the model's generative policy via explicit control tokens (e.g., [RETRIEVE], [ANSWER], [INTERMEDIARY]) without external classifiers.
• 🔄 Self-Triggered Planning: Autonomously decides when to fall back to internal knowledge, how to reformulate targeted queries based on partial reasoning, and when to terminate the search.
• ⚖️ Adaptive Retrieval Depth: Dynamically adjusts the number of retrieval rounds based on question complexity, successfully avoiding redundant searches while extrapolating beyond strict training budgets.
• 🚀 State-of-the-Art Performance: Surpasses strong open-source RAG baselines (e.g., GainRAG, R1-Searcher) and achieves performance competitive with GPT-4o across five QA benchmarks using a much smaller backbone (LLaMA3-8B).
• 🧩 Unified Decoding Trajectory: Tightly couples multi-step reasoning and on-the-fly evidence integration into a single, continuous generation flow.
• 🛠️ Optimized Training Recipe: Employs a structured supervised fine-tuning (SFT) over four distinct behavioral patterns, further refined by rule-based Reinforcement Learning (DAPO) to ensure accurate and balanced retrieval control.

🚀 Quick Start

Installation

git clone https://github.com/WisdomShell/GRIP
cd GRIP
conda create -n GRIP python=3.9
conda activate GRIP
cd GRIP/model/Train
pip install -e .
cd ../
pip install -r requirements.txt

Preparation

Build Wikipedia index

Download the Wikipedia dump.

mkdir wiki_data
cd wiki_data
wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz
gzip -d psgs_w100.tsv.gz

Use Elasticsearch to index the Wikipedia dump

mkdir ret
cd ret
wget -O elasticsearch-7.17.9.tar.gz https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.17.9-linux-x86_64.tar.gz
tar zxvf elasticsearch-7.17.9.tar.gz
rm elasticsearch-7.17.9.tar.gz 
cd elasticsearch-7.17.9
nohup bin/elasticsearch 
python data_generation/index.py --data_path path/to/your/psgs_w100.tsv --index_name wiki

Checkpoints and Datasets

Below are the datasets used for SFT and RL training in our work, and the weights of the already trained GRIP model.

Dataset	HF Dataset Repo
GRIP_SFT_Train_Data	WisdomShell/GRIP_SFT_Data
GRIP_RL_Train_Data	WisdomShell/GRIP_RL_Data

Model	HF Model Repo
Meta-LLaMa-3-8b-GRIP	WisdomShell/LLaMa-3-8b-GRIP

Generation SFT and RL Training Data

Before that, you need to download the NaturalQuestion-open training set, WebQuestions training set and TriviaQA training set, extract their Questions and answers and merge them into a jsonl file,Convert them into the following format:

{
    "question": "",
    "answer":["Answer", ...]
}

Use the Meta-LLaMa-3-8B-Instruct model to execute the following code

bash data_generation/first.sh

Write your OpenAI token into the use_gpt_for_data.py file,and configure C.jsonl file path. After the generation is completed, it will automatically overwrite the original file.

python generation_train_data/use_gpt_for_data.py

Write the directory where A, B, C, and D are located into the merge_dataset.py file. The output path will save SFT_Train_data and RL_Train_data.

python generation_train_data/merge_dataset.py

Train

SFT

1. Data Process

   • Script: Train/examples/data_preprocess/grip/sft.py

   • You need to specify the data_path parameter, indicating the path of the data synthesized by GRIP

     parser.add_argument('--data_path', default='<PATH_TO_RAW_DATASET_ROOT>/SFT_data.jsonl')
     

   • You should specify the name of the dataset for use during subsequent training.

     # The data path is stored in the "datasets" folder by default.
     parser.add_argument('--save_dir', default='datasets/GRIPSFT')
     

2. Train Script

   • Script: Train/examples/sft/run_sft_llama.sh

   • Train using the Base version of the model.

     set -x
     
     NAME=GRIPSFT  # Here to specify the names of the processed training data from the previous step
     torchrun --standalone --nnodes=1 --nproc_per_node=8 -m verl.trainer.fsdp_sft_trainer \
         data.train_files=datasets/$NAME/train.parquet \  
         data.val_files=datasets/$NAME/test.parquet \
         data.prompt_key=extra_info \
         data.response_key=extra_info \
         optim.lr=1e-6 \
         data.prompt_dict_keys=['question'] \
         +data.response_dict_keys=['answer'] \
         data.micro_batch_size=4 \
         model.partial_pretrain=meta-llama/Meta-Llama-3-8B-Base \       #Use Base to Train
         trainer.default_local_dir=/path/to/your/SFT_model \  # Finetuned Model Save Path
         trainer.project_name=GRIPSFT \
         trainer.experiment_name=$NAME \
         trainer.logger=['console'] \  # Report `console` or `wandb`
         trainer.total_epochs=8 \      # Training Epoches
         trainer.default_hdfs_dir=null $@ \
         ulysses_sequence_parallel_size=2 \
         use_remove_padding=true
     

RL

1. Data Process

   • Script: Train/examples/data_preprocess/grip/rl.py

   • You need to specify the data_path parameter, indicating the path of the data synthesized by GRIP

     parser.add_argument('--data_path', default='<PATH_TO_RAW_DATASET_ROOT>/RL_data.jsonl')
     

   • You should specify the name of the dataset for use during subsequent training.

     # The data path is stored in the "datasets" folder by default.
     parser.add_argument('--save_dir', default='datasets/GRIPRL')
     

   • You should specify the name of the data_source for use during subsequent training to select reward model.

     parser.add_argument('--data_source', default='GRIPRL')    # Necessary
     

2. Train Script using DAPO

   • Script: Train/recipe/dapo/dapo_4w_continue_rl_ep3_llama.sh

   • You should modify these parameters to suit RL training.

     ...
      # Paths
      MODEL_PATH=<PATH_TO_SAVE>/GRIPSFT_LLaMa/global_step_xxx  # SFT Checkpoint
      CKPTS_DIR=<PATH_TO_SAVE>/RL_model  # RL Model Save Path
      TRAIN_FILE=datasets/GRIPRL/train.parquet  # RL Datasets
      TEST_FILE=datasets/GRIPRL/test.parquet  # RL Datasets
      ...
     

3. The specific implementation of the Reward Model is in the file Train/verl/utils/reward_score/grip.py.

4. After training, you should merge the slices saved from the model into Hugging Face format by script Train/scripts/merge.sh.

Local Inference using GRIP

Test data format alignment

{
    "question": "Test Query",
    "answer": ["Answer List", ...]
}

Mutil-Turn GRIP Inference

• Main Script: inference/inference.sh

  # Model Saved Path
  parser.add_argument('--model_path', type=str, default="/path/to/your/RL_model/step_xxx")
  # Predicted file output path
  parser.add_argument('--output_file', type=str, default="output/rl_step_xxx_hotpot.jsonl")
  # File  to be predicted 
  parser.add_argument('--input_file', type=str, default="test_data/hotpotQA.jsonl")
  

• This script will generate predicitons by format:

  {
        "Question": "String", 
        "prediction": ["String",......]
    }
  

Eval

python eval/eval.py \
    --references_path test_dataset.jsonl \
    --predictions_path prediction.jsonl

🤝 Contributing

We welcome contributions! See CONTRIBUTING.md for guidelines.

📄 Citation

If you use this repository, please cite the paper.

📝 License

This project is licensed under the Apache 2.0 License - see the LICENSE file for details.

🙏 Acknowledgments

Special thanks to the open-source community and all contributors who made this project possible.