Search-R1 : Apprentissage par renforcement pour former de grands modèles de recherche et de raisonnement

Introduction générale

Search-R1 est un projet open source développé par PeterGriffinJin sur GitHub et construit sur le framework veRL. Il utilise des techniques d'apprentissage par renforcement (RL) pour former de grands modèles de langage (LLM), permettant aux modèles d'apprendre de manière autonome à raisonner et à invoquer des moteurs de recherche pour résoudre des problèmes. Le projet prend en charge des modèles de base tels que Qwen2.5-3B et Llama3.2-3B, et étend le modèle LLM à d'autres modèles. Profondeur de l'eau-R1 et TinyZero. Les utilisateurs peuvent l'utiliser pour former des modèles afin de gérer des tâches à un ou plusieurs tours. Le code, les ensembles de données et les journaux d'expérience sont fournis. Officieldiscuter d'un article ou d'une thèse (ancien)Publié en mars 2025, le modèle et les données du projet peuvent être téléchargés sur le site Hugging Face à l'intention des chercheurs et des développeurs.

Liste des fonctions

• Formation de grands modèles par apprentissage par renforcement pour améliorer le raisonnement et la recherche.
• Prise en charge de l'appel à l'API de Google, Bing, Brave et d'autres moteurs de recherche.
• Fournit des capacités de réglage LoRA et de réglage fin supervisé pour optimiser les performances du modèle.
• Réordonnateur intégré pour améliorer la précision des résultats de recherche.
• Inclut des journaux de laboratoire détaillés et des documents pour appuyer la reproduction des résultats.
• Fonctionnalité de serveur de recherche local permettant de personnaliser facilement les recherches.
• Prise en charge du téléchargement par l'utilisateur d'ensembles de données et de corpus personnalisés.

Utiliser l'aide

Search-R1 s'adresse à des utilisateurs ayant des connaissances de base en programmation et en apprentissage automatique. Vous trouverez ci-dessous un guide d'installation et d'utilisation détaillé pour vous permettre de démarrer rapidement.

Processus d'installation

Pour utiliser Search-R1, vous devez d'abord configurer l'environnement. Les étapes sont les suivantes :

1. Création d'un environnement de recherche-R1
   S'exécute dans le terminal :

conda create -n searchr1 python=3.9
conda activate searchr1

Cela crée un environnement virtuel Python 3.9.

2. Installation de PyTorch
   Installez PyTorch 2.4.0 (compatible avec CUDA 12.1) en entrant la commande suivante :

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

3. Installation de vLLM
   vLLM est une bibliothèque clé pour l'exécution de grands modèles, installer la version 0.6.3 :

pip3 install vllm==0.6.3

Les versions 0.5.4, 0.4.2 ou 0.3.1 sont également disponibles.

4. Installation de veRL
   Exécutez-le dans le répertoire racine du projet :

pip install -e .

Cette opération permet d'installer le framework veRL.

5. Installer les dépendances optionnelles
   Pour améliorer les performances, installez Flash Attention et Wandb :

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

6. Installation de l'environnement Retriever (optionnel)
   Si un serveur d'extraction local est nécessaire, créez un autre environnement :

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

Démarrage rapide

Les étapes suivantes permettent de former un modèle basé sur l'ensemble de données NQ :

1. Télécharger les index et les corpus
   Définissez le chemin d'accès à la sauvegarde et exécutez le programme :

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. Traitement des données NQ
   Exécutez le script pour générer des données de formation :

python scripts/data_process/nq_search.py

3. Démarrer le serveur d'extraction
   Il évolue dans l'environnement du Retriever :

conda activate retriever
bash retrieval_launch.sh

4. Entraînement à la course à pied
   Exécuter dans un environnement Search-R1 :

conda activate searchr1
bash train_ppo.sh

Il utilisera le modèle de base Llama-3.2-3B pour la formation des OPP.

Utilisation d'ensembles de données personnalisés

1. Préparer les données d'assurance qualité
   Les données doivent être au format JSONL et chaque ligne contient les champs suivants :

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

consultation <scripts/data_process/nq_search.py>.

2. Préparation du corpus
   Le corpus doit être au format JSONL, chaque ligne contenant id répondre en chantant contentsComme :

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

   référençable <example/corpus.jsonl>.

3. Corpus d'indexation (optionnel)
   Si vous utilisez la recherche locale, exécutez :

   bash search_r1/search/build_index.sh
   

Appeler un moteur de recherche personnalisé

1. modifications <search_r1/search/retriever_server.py>Les API de configuration.
2. Démarrer le serveur :

   python search_r1/search/retriever_server.py
   

3. Le modèle est transmis par l'intermédiaire de l'outil http://127.0.0.1:8000/retrieve Recherche d'appel.

opération d'inférence

1. Démarrer le serveur d'extraction :

   bash retrieval_launch.sh
   

2. Raisonnement en cours d'exécution :

   python infer.py
   

3. modifications <infer.py> Ligne 7 questionSaisissez la question que vous souhaitez poser.

mise en garde

• L'entraînement nécessite un GPU doté d'au moins 24 Go de mémoire vidéo (par exemple, NVIDIA A100).
• Assurez-vous que la clé API est valide et que la connexion réseau est stable.
• Documents officiels et journaux de laboratoire (<Full experiment log 1> répondre en chantant <Full experiment log 2>) Fournir plus de détails.

Grâce à ces étapes, vous pouvez utiliser Search-R1 pour former un modèle capable de raisonner et d'effectuer des recherches pour traiter une variété de tâches.

scénario d'application

1. expérience de recherche
   Les chercheurs peuvent utiliser Search-R1 pour reproduire les résultats de l'article et explorer l'application de l'apprentissage par renforcement dans la formation des modèles.
2. Développement d'un assistant intelligent
   Les développeurs peuvent former des modèles à intégrer dans les outils de chat pour fournir des capacités de recherche et de raisonnement.
3. Requêtes de connaissances
   Les utilisateurs peuvent l'utiliser pour répondre rapidement à des questions complexes et obtenir des informations actualisées grâce à la recherche.

QA

1. Quels sont les modèles pris en charge par Search-R1 ?
   Actuellement, les modèles de base Qwen2.5-3B et Llama3.2-3B sont pris en charge, les autres modèles doivent être adaptés par vous-même.
2. Quelle est la durée de la formation ?
   En fonction de l'ensemble de données et du matériel, l'entraînement de l'ensemble de données NQ prend environ quelques heures sur un GPU graphique de 24 Go.
3. Comment puis-je vérifier l'efficacité de ma formation ?
   vérifier <Preliminary results> ou consulter le journal de Wandb.