Search-R1: Verstärkungslernen zum Trainieren großer Modelle für Suche und Schlussfolgerungen

Allgemeine Einführung

Search-R1 ist ein Open-Source-Projekt, das von PeterGriffinJin auf GitHub entwickelt wurde und auf dem veRL-Framework aufbaut. Es nutzt Techniken des Reinforcement Learning (RL), um große Sprachmodelle (LLMs) zu trainieren, so dass die Modelle selbstständig lernen, Schlussfolgerungen zu ziehen und Suchmaschinen zur Problemlösung aufzurufen. Das Projekt unterstützt grundlegende Modelle wie Qwen2.5-3B und Llama3.2-3B und erweitert die DeepSeek-R1 und TinyZero-Methoden. Benutzer können damit Modelle für Einzel- oder Mehrrundenaufgaben trainieren, wobei Code, Datensätze und Experimentprotokolle bereitgestellt werden. Offiziellein Papier oder eine These zu diskutieren (alt)Das Projektmodell und die Daten wurden im März 2025 veröffentlicht und stehen Forschern und Entwicklern unter Hugging Face zum Download zur Verfügung.

Funktionsliste

• Training großer Modelle durch Verstärkungslernen zur Verbesserung des Denkens und der Suche.
• Unterstützung für den Aufruf der API von Google, Bing, Brave und anderen Suchmaschinen.
• Bietet die Möglichkeit der LoRA-Abstimmung und der überwachten Feinabstimmung zur Optimierung der Modellleistung.
• Eingebauter, standardmäßiger Reorderer zur Verbesserung der Genauigkeit der Suchergebnisse.
• Enthält detaillierte Laborprotokolle und Unterlagen zur Unterstützung der Reproduktion der Ergebnisse.
• Bietet eine lokale Suchserverfunktionalität zur einfachen Anpassung der Suche.
• Unterstützung für das Hochladen von benutzerdefinierten Datensätzen und Korpora.

Hilfe verwenden

Search-R1 richtet sich an Benutzer mit Grundkenntnissen in Programmierung und maschinellem Lernen. Nachfolgend finden Sie eine detaillierte Installations- und Nutzungsanleitung, damit Sie schnell loslegen können.

Einbauverfahren

Um Search-R1 zu verwenden, müssen Sie zunächst die Umgebung einrichten. Die Schritte sind wie folgt:

1. Erstellen einer Search-R1-Umgebung
   Läuft im Terminal:

conda create -n searchr1 python=3.9
conda activate searchr1

Dadurch wird eine virtuelle Umgebung mit Python 3.9 erstellt.

2. Installation von PyTorch
   Installieren Sie PyTorch 2.4.0 (unterstützt CUDA 12.1) durch Eingabe des folgenden Befehls:

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

3. Installation von vLLM
   vLLM ist eine wichtige Bibliothek für große Modelle, installieren Sie Version 0.6.3:

pip3 install vllm==0.6.3

Die Versionen 0.5.4, 0.4.2 oder 0.3.1 sind ebenfalls verfügbar.

4. Installation von veRL
   Führen Sie es im Stammverzeichnis des Projekts aus:

pip install -e .

Dadurch wird das veRL-Framework installiert.

5. Optionale Abhängigkeiten installieren
   Um die Leistung zu verbessern, installieren Sie Flash Attention und Wandb:

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

6. Installation der Retriever-Umgebung (optional)
   Wenn ein lokaler Abrufserver erforderlich ist, erstellen Sie eine weitere Umgebung:

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

Schnellstart

Im Folgenden werden die Schritte zum Trainieren eines Modells auf der Grundlage des NQ-Datensatzes beschrieben:

1. Indizes und Korpora herunterladen
   Legen Sie den Speicherpfad fest und führen Sie ihn aus:

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. Verarbeitung von NQ-Daten
   Führen Sie das Skript aus, um Trainingsdaten zu erzeugen:

python scripts/data_process/nq_search.py

3. Starten Sie den Abrufserver
   Läuft in der Retriever-Umgebung:

conda activate retriever
bash retrieval_launch.sh

4. RL-Training durchführen
   In einer Search-R1-Umgebung ausführen:

conda activate searchr1
bash train_ppo.sh

Dabei wird das Basismodell Llama-3.2-3B für die PPO-Ausbildung verwendet.

Benutzerdefinierte Datensätze verwenden

1. Aufbereitung von QA-Daten
   Die Daten müssen im JSONL-Format vorliegen und jede Zeile muss die folgenden Felder enthalten:

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

Beratung <scripts/data_process/nq_search.py>.

2. Aufbereitung des Korpus
   Der Korpus muss im JSONL-Format vorliegen, wobei jede Zeile Folgendes enthält id im Gesang antworten contentsAls:

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

   verweisbar <example/corpus.jsonl>.

3. Indexierungskorpus (optional)
   Wenn Sie die lokale Suche verwenden, führen Sie aus:

   bash search_r1/search/build_index.sh
   

Aufrufen einer benutzerdefinierten Suchmaschine

1. Änderungen <search_r1/search/retriever_server.py>Die Konfigurations-APIs.
2. Starten Sie den Server:

   python search_r1/search/retriever_server.py
   

3. Das Modell wird durch den http://127.0.0.1:8000/retrieve Suche aufrufen.

Inferenzoperation

1. Starten Sie den Abrufserver:

   bash retrieval_launch.sh
   

2. Running Reasoning:

   python infer.py
   

3. Änderungen <infer.py> Zeile 7 questiongeben Sie die Frage ein, die Sie stellen möchten.

caveat

• Das Training erfordert einen Grafikprozessor mit mindestens 24 GB Videospeicher (z. B. NVIDIA A100).
• Vergewissern Sie sich, dass der API-Schlüssel gültig ist und die Netzwerkverbindung stabil ist.
• Offizielle Papiere und Laborjournale (<Full experiment log 1> im Gesang antworten <Full experiment log 2>) Geben Sie weitere Einzelheiten an.

Mit diesen Schritten können Sie Search-R1 verwenden, um ein Modell zu trainieren, das Schlussfolgerungen ziehen und suchen kann, um eine Vielzahl von Aufgaben zu bewältigen.

Anwendungsszenario

1. Forschungsexperiment
   Forscher können Search-R1 verwenden, um die Ergebnisse des Artikels zu reproduzieren und die Anwendung des Verstärkungslernens beim Modelltraining zu untersuchen.
2. Entwicklung intelligenter Assistenten
   Die Entwickler können Modelle trainieren, die in Chat-Tools integriert werden können, um Such- und Schlussfolgerungsfunktionen bereitzustellen.
3. Wissensabfragen
   Die Nutzer können damit schnell komplexe Fragen beantworten und über die Suche aktuelle Informationen erhalten.

QA

1. Welche Modelle werden von Search-R1 unterstützt?
   Derzeit werden die Basismodelle Qwen2.5-3B und Llama3.2-3B unterstützt, andere Modelle müssen selbst angepasst werden.
2. Wie lange dauert die Ausbildung?
   Je nach Datensatz und Hardware dauert das Training des NQ-Datensatzes auf einem 24-GB-Grafikprozessor etwa einige Stunden.
3. Wie überprüfe ich die Wirksamkeit meiner Ausbildung?
   auschecken <Preliminary results> oder prüfen Sie das Wandb-Protokoll.