性能概览：

什么是s3？

s3框架

s3 是一个简单而强大的框架，用于在检索增强生成（RAG）中训练搜索代理。它教会语言模型如何更有效地进行搜索——而无需更改生成器本身。通过仅关注搜索组件，s3 仅使用先前方法所需数据的一小部分，便能在问答任务上取得优异的性能。它具有模块化、高效的特点，并且可以与任何黑盒大语言模型无缝协作。

目录

• 📦 安装
• 💡 准备工作
• 🏋️ 运行训练
• 🔍 运行搜索/检索
• 📈 运行评估

📦 安装

搜索器与生成器环境

conda create -n s3 python=3.9
# 安装 torch [或者你可以跳过这一步，让 vllm 自动为你安装合适的版本]
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
# 安装 vllm
pip3 install vllm==0.6.3 # 或者你也可以安装 0.5.4、0.4.2 和 0.3.1
pip3 install ray

# verl
# cd code
pip install -e .

# flash attention 2
pip3 install flash-attn --no-build-isolation

# 我们使用 pyserini 来进行高效的检索和评估
pip install pyserini    # 我们使用的版本是 0.22.1

# 提升开发体验
pip install wandb IPython matplotlib huggingface_hub

检索器环境

conda create -n ret python=3.10
conda activate ret

conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install transformers datasets pyserini
conda install -c pytorch -c nvidia faiss-gpu=1.8.0
pip install uvicorn fastapi

💡 准备工作

下载索引与语料库

python scripts/download.py --save_path $save_path
cat $save_path/part_* > $save_path/e5_Flat.index
gzip -d $save_path/wiki-18.jsonl.gz

预计算朴素 RAG 初始化（你也可以在这里下载我们处理好的数据：huggingface）

# 部署检索器
bash scripts/deploy_retriever/retrieval_launch.sh # 或者对于 MedCorp 语料库，可以使用 scripts/deploy_retriever/retrieval_launch_mirage.sh。
# 部署生成器
bash generator_llms/host.sh # 将 tensor-parallel-size 修改为你所使用的 GPU 数量
# 运行预计算
bash scripts/precompute.sh # 这一步需要一些时间，因为它会为训练预先计算朴素 RAG 缓存

🏋️ 运行训练

此步骤用于 s3 的训练

# 部署检索器
bash scripts/deploy_retriever/retrieval_launch.sh 
# 部署生成器
bash generator_llms/host.sh
# 运行训练
bash scripts/train/train_s3.sh

🔍 运行搜索/检索

此步骤用于收集 s3 及基线模型的上下文

s3

# 部署检索器
bash scripts/deploy_retriever/retrieval_launch.sh 
# 运行 s3 推理
bash scripts/s3_inference/evaluate-8-3-3.sh

bash scripts/deploy_retriever/retrieval_launch.sh # 或 retrieval_launch_bm25.sh # 部署检索器
bash scripts/baselines/rag.sh # 运行 RAG 

bash retrieval_launch_bm25.sh # 部署 BM25 模型
bash generator_llms/deepretrieval.sh # 部署 DeepRetrieval 模型
bash scripts/baselines/deepretrieval.sh # 运行 DeepRetrieval 查询改写 + 检索

bash retrieval_launch.sh # 部署 e5 检索器
bash scripts/baselines/search_r1.sh # 运行 Search-R1

bash retrieval_launch.sh # 部署 e5 检索器
python scripts/baselines/ircot.py

bash retrieval_launch.sh # 部署 e5 检索器
bash scripts/baselines/search_o1.sh # 运行 Search-o1

📈 运行评估

此步骤用于评估 s3 及基线模型

bash scripts/evaluation/run.sh

问答

自定义数据？

如果你想在自己的语料库或数据集上测试 s3，可以参考这个提交，了解如何构建你自己的流水线：commit 8420538

如何复现结果？

目前已有几位开发者成功复现了我们的实验结果。如果您有任何疑问或遇到问题，欢迎随时提交 issue——我们很乐意提供一对一的指导（请参阅这个示例)。

尽管自行复现该模型非常简单——我们实际上建议从头开始训练，因为评估通常比训练耗时得多——我们仍提供了一个参考检查点：s3-8-3-3-20steps，该模型仅用约一小时便完成了训练。

引用

@article{jiang2025s3,
  title={s3：通过强化学习训练搜索智能体并不需要大量数据},
  author={Jiang, Pengcheng and Xu, Xueqiang and Lin, Jiacheng and Xiao, Jinfeng and Wang, Zifeng and Sun, Jimeng and Han, Jiawei},
  journal={arXiv 预印本 arXiv:2505.14146},
  year={2025}
}

感谢您对我们工作的关注！

s3 快速上手指南

s3 是一个高效且轻量级的框架，旨在通过强化学习（RL）训练搜索代理（Search Agent）。它专注于优化检索组件，无需修改生成模型本身，即可在问答任务中实现卓越性能，且所需训练数据量远少于传统方法。

环境准备

本项目需要两个独立的 Conda 环境：一个用于搜索器与生成器（Searcher & Generator），另一个用于检索器（Retriever）。

系统要求：

• Python 3.9 (搜索/生成环境) / Python 3.10 (检索环境)
• NVIDIA GPU (推荐 CUDA 12.1)
• 足够的显存以运行 vLLM 和检索模型

前置依赖：

• Conda
• Git

安装步骤

1. 配置搜索器与生成器环境

# 创建环境
conda create -n s3 python=3.9
conda activate s3

# 安装 PyTorch (CUDA 12.1)
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121

# 安装 vLLM 及相关依赖
pip3 install vllm==0.6.3
pip3 install ray

# 安装 s3 核心代码 (假设已克隆仓库并进入 code 目录)
pip install -e .

# 安装 Flash Attention 2
pip3 install flash-attn --no-build-isolation

# 安装检索评估工具 Pyserini
pip install pyserini

# 安装辅助工具
pip install wandb IPython matplotlib huggingface_hub

2. 配置检索器环境

# 创建独立环境
conda create -n ret python=3.10
conda activate ret

# 安装 PyTorch 全家桶 (CUDA 12.1)
conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia

# 安装核心库
pip install transformers datasets pyserini

# 安装 FAISS GPU 版本
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

# 安装 API 服务依赖
pip install uvicorn fastapi

基本使用

以下是从零开始运行 s3 的核心流程，包括数据准备、训练及推理。

第一步：数据准备

下载索引和语料库，并预处理初始化的 RAG 缓存。

# 下载数据 (替换 $save_path 为你的保存路径)
python scripts/download.py --save_path $save_path

# 合并索引文件
cat $save_path/part_* > $save_path/e5_Flat.index

# 解压语料库
gzip -d $save_path/wiki-18.jsonl.gz

# 部署检索器和生成器服务
bash scripts/deploy_retriever/retrieval_launch.sh
bash generator_llms/host.sh 
# 注意：请根据显卡数量修改 host.sh 中的 tensor-parallel-size 参数

# 预计算朴素 RAG 初始化缓存 (此步骤耗时较长)
bash scripts/precompute.sh

> 提示：如果不想自行预处理，可直接从 HuggingFace 下载处理好的数据：pat-jj/s3_processed_data

第二步：运行训练

启动强化学习训练流程。需确保检索器和生成器服务已在后台运行。

# 确保服务已启动 (若第一步已完成可跳过)
bash scripts/deploy_retriever/retrieval_launch.sh
bash generator_llms/host.sh

# 开始训练 s3
bash scripts/train/train_s3.sh

第三步：运行搜索/推理

使用训练好的模型进行检索测试。

# 确保检索器服务运行中
bash scripts/deploy_retriever/retrieval_launch.sh

# 执行 s3 推理评估
bash scripts/s3_inference/evaluate-8-3-3.sh

第四步：结果评估

对搜索结果进行最终评分。

bash scripts/evaluation/run.sh