DeepRetrieval - 通过强化学习利用大语言模型攻破真实搜索引擎与检索系统

      
⭐️ 请给我们的仓库点个星，以便及时获取最新功能更新！！

最新消息

• [2025-05-19] 如果您对 RAG 任务特别感兴趣，我们推荐您查看我们的最新工作 s3，这是一个专为 RAG 设计的高效训练框架。
• [2025-05-06] 您现在可以轻松部署我们的模型并调用 API 进行查询改写。更多详情请参见 这里。
• [2025-02-28] 🎉 我们很高兴地发布了 DeepRetrieval 的代码！

什么是 DeepRetrieval？

DeepRetrieval 是一种新颖的强化学习方法，用于训练大型语言模型（LLMs）生成查询，以提升信息检索性能。与依赖于标注查询增强对的监督学习的传统方法不同，DeepRetrieval 让模型通过直接的试错过程进行学习，并以检索指标作为奖励，从而生成能够最大化检索效果的查询。

该系统的工作原理是：让 LLM 在 <think> 部分生成推理步骤，随后在 <answer> 部分给出最终的增强查询。这种结构化的方法能够在确定查询形式之前进行明确的思维链式推理。

核心特性与成果

• 无需监督：无需昂贵的人工标注或蒸馏得到的参考查询
• 强大性能：显著超越以往的最先进方法
  • 文献搜索的召回率为 65.07%（而之前的 SOTA 为 24.68%）
  • 临床试验搜索的召回率为 63.18%（而之前的 SOTA 为 32.11%）
• 应用广泛：适用于多种检索任务：(1) 使用真实世界搜索引擎的文献检索 (2) 证据检索 (3) 经典的信息检索 (4) SQL 数据库检索
• 参数高效：仅使用 3B 参数就取得了优异的效果，优于更大的模型如 GPT-4o 和 Claude-3.5-Sonnet

PubMed 搜索引擎上的 Wandb 训练日志示例

⭐️ 请给我们的仓库点个星，以便随时了解激动人心的新功能和改进！🌟

目录

• 📦 安装
• ⚡️ 易用的查询改写 API
• 🫧 开始使用
• 🏃 运行训练
• 🧐 运行评估
• 📚 引用 DeepRetrieval

📦 安装

通用安装（适用于所有检索方法）：

conda create -n zero 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
# 提升开发体验
pip install wandb IPython matplotlib huggingface_hub

(如果您只想运行搜索引擎检索，则可以跳过以下步骤。)

如果使用稀疏检索（例如 BM25）或密集检索（例如 DPR），请同时安装以下内容：

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

# 如果您尚未安装 faiss，请安装：
pip install faiss-gpu==1.7.2    # 我们使用的版本是 1.7.2

# 如果您尚未安装 Java，请安装：
pip install install-jdk && python -c "import jdk; jdk.install('11')"

# 支持 SQL 执行
pip install func_timeout

⚡️ 易用的查询改写 API

sh vllm_host.sh # 您可以指定本地端口和要使用的模型
python query_rewrite.py --query "谁在 2025 年构建了 DeepRetrieval？"
# 原始查询：谁在 2025 年构建了 DeepRetrieval？
# 改写后的查询：（DeepRetrieval 系统，它是在 2025 年构建的）

🫧 开始使用

1. 数据下载/预处理

我们提供了两种数据准备方式：

cd code
# 列出可用的数据集
python download_datasets.py --list_only --repo_id DeepRetrieval/datasets

# 下载所有数据集
python download_datasets.py --repo_id DeepRetrieval/datasets --output_dir ./data

# 或者下载特定类别/数据集
python download_datasets.py --categories search_engine --datasets pubmed_32 --output_dir ./data

cd code
python download_datasets.py --categories raw_data --datasets pubmed --output_dir ./data
python data_preprocess/pubmed_32.py

2. 获取您的搜索引擎 API 密钥（若使用搜索引擎则必需）

以 PubMed 为例，您可以按照 此处 的说明获取密钥。（PubMed API 是 100% ✨免费✨使用的！）

然后将其放入 code/verl/utils/reward_score/apis/ 目录下，命名为 pubmed_api.key。

3. 奖励函数相关（可选）

奖励设计（例如在 code/verl/utils/reward_score/pubmed.py 中）：

4. 自定义监控信息（可选）

修改 code/verl/trainer/ppo/ray_trainer.py 中的 compute_reward_metrics()

🏃 运行训练

conda activate zero

以 PubMed 为例：

sh scripts/train/pubmed_32.sh 

（如果出现显存不足的情况，可以尝试在脚本中添加 critic.model.enable_gradient_checkpointing=True）

训练期间的思维长度与查询长度（PubMed）

🧐 运行评估

sh scripts/eval/pubmed_32.sh

（您无需进行训练即可运行此脚本，因为它会默认从Huggingface下载我们训练好的模型）

搜索引擎上的结果

*表：不同模型和方法在文献搜索及临床试验搜索任务上的对比。

注： LEADS-7B 是一款最先进的文献挖掘大模型，基于2万篇综述和40万篇文献进行训练 [https://arxiv.org/pdf/2501.16255]

⚠️ 常见问题

⚙️ 集群设置：CUDA与Conda环境

# --- 加载CUDA模块 ---
module load cuda/12.6

# --- 激活Conda环境 ---
source ~/miniconda/etc/profile.d/conda.sh  # 如果您的Conda安装路径不同，请相应调整
conda activate zero

ERROR: pip的依赖解析器目前未考虑到所有已安装的包。
这种行为是以下依赖冲突的根源。
matplotlib 3.9.4需要pyparsing>=2.3.1，但该包未安装。
mne 1.8.0需要decorator，但该包未安装。
pyhealth 1.1.4需要bitsandbytes，但该包未安装。

pip install pyparsing decorator bitsandbytes

pyhealth 1.1.4需要numpy<2.0，但您当前安装的是numpy 2.0.2，两者不兼容。
outlines 0.0.46需要numpy<2.0.0，但您当前安装的是numpy 2.0.2，同样不兼容。
vllm 0.6.3需要numpy<2.0.0，但您当前安装的是numpy 2.0.2，无法兼容。

pip install "numpy<2.0.0" --force-reinstall
pip install transformers==4.46.3
pip install vllm==0.6.3
pip install verl==0.1
pip install outlines==0.0.46
pip install pyserini

🤝 致谢

本实现主要基于verl和PySerini。实验中的基础模型为Qwen2.5-3B-Instruct。我们衷心感谢他们对开源社区的贡献。

📚 引用DeepRetrieval

@article{jiang2025deepretrieval,
  title={Deepretrieval: Hacking real search engines and retrievers with large language models via reinforcement learning},
  author={Jiang, Pengcheng and Lin, Jiacheng and Cao, Lang and Tian, Runchu and Kang, SeongKu and Wang, Zifeng and Sun, Jimeng and Han, Jiawei},
  journal={arXiv preprint arXiv:2503.00223},
  year={2025}
}

感谢您的关注！😊

DeepRetrieval 快速上手指南

DeepRetrieval 是一个基于强化学习（RL）的创新框架，旨在训练大语言模型（LLM）生成更高效的检索查询。它无需人工标注数据，通过“试错”机制直接利用检索指标作为奖励，显著提升在文献搜索、临床试验搜索等任务中的召回率。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: 3.9
• GPU: 支持 CUDA 的 NVIDIA 显卡（建议显存充足以运行 vLLM 和训练）
• 其他依赖:
  • Java (如果使用稀疏检索如 BM25)
  • Conda (用于环境管理)

安装步骤

1. 创建并激活 Conda 环境

conda create -n zero python=3.9
conda activate zero

2. 安装核心依赖

首先安装 PyTorch（根据您的需求选择 CUDA 版本，此处以 cu121 为例），然后安装 vLLM 和其他核心库。

# 安装 torch (若跳过此步，vllm 可能会自动安装对应版本)
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121

# 安装 vllm 和 ray
pip3 install vllm==0.6.3
pip3 install ray

# 安装 verl 框架
cd code
pip install -e .

# 安装 flash-attn (加速注意力机制)
pip3 install flash-attn --no-build-isolation

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

3. 安装可选依赖（针对特定检索任务）

如果您计划使用稀疏检索（如 BM25）、稠密检索（如 DPR）或 SQL 搜索，请额外安装以下包：

# 安装 pyserini (用于高效检索和评估)
pip install pyserini

# 安装 faiss-gpu (用于向量检索)
pip install faiss-gpu==1.7.2

# 安装 Java (pyserini 依赖)
pip install install-jdk && python -c "import jdk; jdk.install('11')"

# 支持 SQL 执行
pip install func_timeout

基本使用

DeepRetrieval 提供了简单的 API 用于查询重写（Query Rewriting），这是其最核心的功能之一。

1. 启动本地服务

运行以下脚本启动 vLLM 服务。您可以指定端口和使用的模型。

sh vllm_host.sh

2. 执行查询重写

使用提供的 Python 脚本发送原始查询，模型将返回优化后的查询语句。

python query_rewrite.py --query "Who built DeepRetrieval in 2025?"

预期输出示例：

# Original query: Who built DeepRetrieval in 2025?
# Rewritten query: (The DeepRetrieval system, which was built in 2025)

3. (进阶) 运行评估

如果您想直接体验训练好的模型效果（无需自行训练），可以运行评估脚本。默认情况下，它会自动从 Hugging Face 下载预训练模型。

sh scripts/eval/pubmed_32.sh

提示：若要使用真实搜索引擎（如 PubMed）进行强化学习训练，您需要在 code/verl/utils/reward_score/apis/ 目录下配置相应的 API Key（例如 pubmed_api.key）。