Search-R1: 検索と推論のための大規模モデルを学習する強化学習

はじめに

Search-R1はPeterGriffinJinによってGitHub上で開発されたオープンソースプロジェクトで、veRLフレームワーク上に構築されている。強化学習(RL)のテクニックを使って大規模言語モデル(LLM)を学習し、モデルが自律的に推論を学習し、問題を解決するために検索エンジンを呼び出すことを可能にする。このプロジェクトはQwen2.5-3BやLlama3.2-3Bのような基本的なモデルをサポートし、以下のような拡張を行います。 ディープシーク-R1 とTinyZeroメソッドで構成されている。コード、データセット、実験ログが提供されている。公式論文について論じる2025年3月に発表されたこのプロジェクト・モデルとデータは、研究者や開発者向けにHugging Faceでダウンロードできる。

機能一覧

• 強化学習によって大規模なモデルを訓練し、推論と探索を改善する。
• Google、Bing、Brave、その他の検索エンジンのAPI呼び出しに対応。
• モデルのパフォーマンスを最適化するために、LoRAチューニングと監視された微調整機能を提供します。
• 検索結果の精度を向上させる、既製のリフォーマー内蔵。
• 結果の再現をサポートする詳細な実験記録と論文を含む。
• 検索を簡単にカスタマイズできるローカル検索サーバー機能を提供。
• カスタマイズされたデータセットやコーパスのユーザーによるアップロードをサポート。

ヘルプの使用

Search-R1は、プログラミングと機械学習の基礎知識を持つユーザーを対象としています。以下は、すぐに使い始められるよう、詳しいインストールと使用方法のガイドです。

設置プロセス

Search-R1を使うには、まず環境を整える必要がある。手順は以下の通り：

1. Search-R1環境の構築
   ターミナルで実行：

conda create -n searchr1 python=3.9
conda activate searchr1

これはPython 3.9の仮想環境を作成する。

2. PyTorchのインストール
   以下のコマンドを入力して、PyTorch 2.4.0（CUDA12.1をサポート）をインストールします：

pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121

3. vLLMのインストール
   ブイエルエルエム は大規模なモデルを実行するための重要なライブラリで、バージョン0.6.3をインストールしてください：

pip3 install vllm==0.6.3

バージョン0.5.4、0.4.2、0.3.1もあります。

4. veRLのインストール
   プロジェクトのルート・ディレクトリで実行する：

pip install -e .

これでveRLフレームワークがインストールされる。

5. オプションの依存関係をインストールする
   パフォーマンスを向上させるには、Flash AttentionとWandbをインストールしてください：

pip3 install flash-attn --no-build-isolation
pip install wandb

6. レトリーバー環境のインストール（オプション）
   ローカル検索サーバーが必要な場合は、別の環境を作成する：

conda create -n retriever python=3.10
conda activate retriever
conda install pytorch==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

クイックスタート

以下は、NQデータセットに基づいてモデルを学習する手順である：

1. 索引とコーパスのダウンロード
   保存パスを設定して実行する：

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

2. NQデータの処理
   スクリプトを実行してトレーニングデータを生成する：

python scripts/data_process/nq_search.py

3. 検索サーバーの起動
   レトリーバー環境で走る：

conda activate retriever
bash retrieval_launch.sh

4. ランRLトレーニング
   Search-R1環境で実行する：

conda activate searchr1
bash train_ppo.sh

PPOトレーニングにはLlama-3.2-3Bのベースモデルを使用する。

カスタムデータセットの使用

1. QAデータの準備
   データはJSONL形式である必要があり、各行には以下のフィールドが含まれる：

{
"data_source": "web",
"prompt": [{"role": "user", "content": "问题"}],
"ability": "fact-reasoning",
"reward_model": {"style": "rule", "ground_truth": "答案"},
"extra_info": {"split": "train", "index": 1}
}

協議 <scripts/data_process/nq_search.py>.

2. コーパスの準備
   コーパスはJSONL形式である必要があります。 id 歌で応える contents例えば

   {"id": "0", "contents": "文本内容"}
   

   参照可能 <example/corpus.jsonl>.

3. コーパスの索引付け（オプション）
   ローカル検索を使用する場合は、実行する：

   bash search_r1/search/build_index.sh
   

カスタム検索エンジンの呼び出し

1. 修正 <search_r1/search/retriever_server.py>コンフィギュレーションAPI。
2. サーバーを起動します：

   python search_r1/search/retriever_server.py
   

3. モデルは http://127.0.0.1:8000/retrieve 電話検索

推論演算

1. 検索サーバーを起動する：

   bash retrieval_launch.sh
   

2. 走る推理：

   python infer.py
   

3. 修正 <infer.py> 7行目 question質問したい内容を入力してください。

ほら

• トレーニングには、少なくとも24GBのビデオメモリを搭載したGPU（NVIDIA A100など）が必要です。
• APIキーが有効で、ネットワーク接続が安定していることを確認する。
• 公式論文とラボジャーナル(<Full experiment log 1> 歌で応える <Full experiment log 2>)詳細を説明してください。

以上の手順で、Search-R1を使って推論と検索ができるモデルを訓練し、さまざまなタスクに対応できるようになる。

アプリケーションシナリオ

1. 研究実験
   研究者はSearch-R1を使って論文の結果を再現し、モデル学習における強化学習の応用を探ることができる。
2. インテリジェント・アシスタント開発
   開発者は、チャットツールに統合されるモデルをトレーニングし、検索や推論機能を提供することができる。
3. ナレッジ・クエリー
   ユーザーは、複雑な質問に素早く答えたり、検索を通じて最新の情報を得たりするために使うことができる。

品質保証

1. Search-R1の対応機種は？
   現在、Qwen2.5-3BとLlama3.2-3Bがサポートされています。
2. トレーニングにはどのくらい時間がかかりますか？
   データセットとハードウェアにもよるが、NQデータセットは24GBのグラフィックスGPUで学習に数時間かかる。
3. トレーニングの効果を確認するには？
   調べる <Preliminary results> またはWandbのログを確認してください。