Search-R1 使用与部署指南

1. 项目介绍

Search-R1 是一个基于 veRL 的强化学习训练框架，旨在训练能够进行推理和调用搜索引擎的大型语言模型（LLM）。通过规则基础的奖励机制（RL），3B 基础模型（如 Qwen2.5-3b-base 和 Llama3.2-3b-base）能够自主发展推理和搜索引擎调用能力。

2. 项目快速启动

环境准备

首先，创建并激活 Python 虚拟环境：

conda create -n searchr1 python=3.9
conda activate searchr1

安装必要的依赖：

pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
pip3 install vllm==0.6.3
pip3 install flash-attn --no-build-isolation
pip install wandb

如果需要使用本地检索器作为搜索引擎，还需要安装以下环境：

conda create -n retriever python=3.10
conda activate retriever
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
conda install -c pytorch -c nvidia faiss-gpu=1.8.0
pip install uvicorn fastapi

数据准备

下载并解压索引和语料库：

save_path=/the/path/to/save
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

处理 NQ 数据集：

python scripts/data_process/nq_search.py

启动检索服务

conda activate retriever
bash retrieval_launch.sh

模型训练

conda activate searchr1
bash train_ppo.sh

3. 应用案例和最佳实践

推理使用

启动本地检索服务后，可以进行推理：

conda activate searchr1
python infer.py

你可以修改 infer.py 文件中第 7 行的 question 变量，将其设置为你感兴趣的问题。

自定义数据集

你可以使用自己的数据集进行训练。数据集应该是一个包含问题、答案和相关信息的字典。可以参考 scripts/data_process/nq_search.py 文件中的数据处理示例。

自定义搜索引擎

可以启动一个本地或远程的搜索引擎服务，并让 LLM 通过调用搜索 API 来使用它。可以参考 search_r1/search/retriever_server.py 文件来启动一个本地检索器服务。

4. 典型生态项目

Search-R1 可以与多种类型的开源项目配合使用，例如：

• 检索器项目，如 Elasticsearch、Faiss 等。
• 自然语言处理工具，如 Transformers、SpaCy 等。
• 强化学习库，如 Stable Baselines、Ray(RLlib) 等。

通过整合这些项目，可以构建更加强大和灵活的搜索和推理系统。