Stop-RAG: Value-Based Retrieval Control for Iterative RAG

This is the official repository for our paper **"Stop-RAG: Value-Based Retrieval Control for Iterative RAG"**. ## Setup We recommend using [uv](https://docs.astral.sh/uv/getting-started/installation/) with Python 3.11 for dependency management. To install the necessary dependencies, run: ```bash pip install uv uv sync --no-dev ``` ## Replication To run Stop-RAG training and evaluate the results, follow the steps below. All experiments were run on 4 H100 GPUs. There are three types of variables to set: - `DATASET`: the name of the dataset to use (`musique` or `hotpotqa` or `2wikimultihopqa`) - `RETRIEVER`: the type of retriever to use (`contriever` or `bm25`) - `METHOD`: the base pipeline to use (`ours` or `corag`) ### 1. Download datasets First, download all datasets and build retrieval embeddings from the corpora. ```bash ./scripts/download.sh {RETRIEVER} ``` Datasets will be stored in `data/raw/{DATASET}`, and the retrieval corpora, embeddings and indexes will be saved in `data/corpus/{DATASET}`. ### 2. Prepare training and evaluation data Run the chosen pipeline to prepare the training and evaluation data. ```bash ./scripts/dataset.sh {DATASET} {RETRIEVER} {METHOD} ``` The processed data will be saved in `data/processed/{DATASET}/{METHOD}/{RETRIEVER}`. ### 3. Run Stop-RAG training Now run the training script. WandB logging is enabled by default, so the environment variables `WANDB_API_KEY` and `WANDB_ENTITY` must be set. To disable WandB logging, set `WANDB_DISABLED=true`. ```bash ./scripts/train.sh {DATASET} {RETRIEVER} {METHOD} ``` ### 4. Run evaluation To run evaluation, first compute the scores for all trained checkpoints and find the best checkpoint and threshold. ```bash ./scripts/stop_rag_find_best.sh {DATASET} {RETRIEVER} {METHOD} ``` This script will print the best checkpoint and threshold. Substitute these values into the following command to run the final evaluation. ```bash ./scripts/stop_rag_test.sh {DATASET} {RETRIEVER} {METHOD} {CKPT} {THRESHOLD} ``` ### (Optional) 5. Evaluate LLM-Stop You can also evaluate the **LLM-Stop** baseline for comparison. ```bash ./scripts/llm_stop_test.sh {DATASET} {RETRIEVER} {METHOD} ``` ## Acknowledgments The code for Contriever is adapted from [EfficientRAG](https://github.com/NIL-zhuang/EfficientRAG-official), and the dataset download scripts are adapted from [IRCoT](https://github.com/StonyBrookNLP/ircot).