Search-R1: Aprendizaje por refuerzo para entrenar grandes modelos de búsqueda y razonamiento

Introducción general

Search-R1 es un proyecto de código abierto desarrollado por PeterGriffinJin en GitHub y basado en el framework veRL. Utiliza técnicas de aprendizaje por refuerzo (RL) para entrenar grandes modelos lingüísticos (LLM), lo que permite a los modelos aprender de forma autónoma a razonar e invocar motores de búsqueda para resolver problemas. El proyecto es compatible con modelos básicos como Qwen2.5-3B y Llama3.2-3B, y amplía los modelos de DeepSeek-R1 y TinyZero. Los usuarios pueden utilizarlo para entrenar modelos que gestionen tareas de una o varias rondas, con código, conjuntos de datos y registros de experimentos. Oficialdebatir sobre una ponencia o tesis (antiguo)El modelo y los datos del proyecto, publicados en marzo de 2025, pueden descargarse en Hugging Face para investigadores y desarrolladores.

Lista de funciones

• Entrenamiento de grandes modelos mediante aprendizaje por refuerzo para mejorar el razonamiento y la búsqueda.
• Soporte para llamar a Google, Bing, Brave y otros motores de búsqueda API.
• Proporciona capacidades de ajuste LoRA y de ajuste fino supervisado para optimizar el rendimiento del modelo.
• Reordenador incorporado para mejorar la precisión de los resultados de búsqueda.
• Incluye registros de laboratorio detallados y documentos para apoyar la reproducción de los resultados.
• Proporciona funciones de servidor de búsqueda local para facilitar la personalización de las búsquedas.
• Posibilidad de que el usuario cargue conjuntos de datos y corpus personalizados.

Utilizar la ayuda

Search-R1 está dirigido a usuarios con conocimientos básicos de programación y aprendizaje automático. A continuación encontrará una guía detallada de instalación y uso para empezar a utilizarlo rápidamente.

Proceso de instalación

Para utilizar Search-R1, primero hay que configurar el entorno. Los pasos son los siguientes:

1. Creación de un entorno Search-R1
   Se ejecuta en el terminal:

conda create -n searchr1 python=3.9
conda activate searchr1

Esto crea un entorno virtual Python 3.9.

2. Instalación de PyTorch
   Instale PyTorch 2.4.0 (soporta CUDA 12.1) introduciendo el siguiente comando:

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

3. Instalación de vLLM
   vLLM es una biblioteca clave para ejecutar grandes modelos, instale la versión 0.6.3:

pip3 install vllm==0.6.3

También están disponibles las versiones 0.5.4, 0.4.2 o 0.3.1.

4. Instalación de veRL
   Ejecútelo en el directorio raíz del proyecto:

pip install -e .

Esto instalará el framework veRL.

5. Instalar dependencias opcionales
   Para mejorar el rendimiento, instala Flash Attention y Wandb:

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

6. Instalación de Retriever Environment (opcional)
   Si se necesita un servidor de recuperación local, cree otro entorno:

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

Inicio rápido

A continuación se indican los pasos para entrenar un modelo basado en el conjunto de datos NQ:

1. Descargar índices y corpus
   Establece la ruta de guardado y ejecútalo:

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. Tratamiento de datos NQ
   Ejecute el script para generar datos de entrenamiento:

python scripts/data_process/nq_search.py

3. Iniciar el servidor de recuperación
   Se ejecuta en el entorno Retriever:

conda activate retriever
bash retrieval_launch.sh

4. Entrenamiento Run RL
   Ejecutar en un entorno Search-R1:

conda activate searchr1
bash train_ppo.sh

Se utilizará el modelo base Llama-3.2-3B para la formación PPO.

Utilización de conjuntos de datos personalizados

1. Preparar los datos de control de calidad
   Los datos deben estar en formato JSONL y cada fila debe contener los siguientes campos:

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

consulta <scripts/data_process/nq_search.py>.

2. Preparación del corpus
   El corpus debe estar en formato JSONL, y cada línea debe contener id responder cantando contentsComo:

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

   referenciable <example/corpus.jsonl>.

3. Indexación del corpus (opcional)
   Si utiliza la búsqueda local, ejecútela:

   bash search_r1/search/build_index.sh
   

Llamar a un motor de búsqueda personalizado

1. modificaciones <search_r1/search/retriever_server.py>Las API de configuración.
2. Inicie el servidor:

   python search_r1/search/retriever_server.py
   

3. El modelo pasa por el http://127.0.0.1:8000/retrieve Búsqueda de llamadas.

operación de inferencia

1. Inicie el servidor de recuperación:

   bash retrieval_launch.sh
   

2. Razonamiento de carrera:

   python infer.py
   

3. modificaciones <infer.py> Línea 7 questionintroduzca la pregunta que desea formular.

advertencia

• El entrenamiento requiere una GPU con al menos 24 GB de memoria de vídeo (por ejemplo, NVIDIA A100).
• Asegúrese de que la clave API es válida y de que la conexión de red es estable.
• Documentos oficiales y revistas de laboratorio (<Full experiment log 1> responder cantando <Full experiment log 2>) Proporcione más detalles.

Con estos pasos, puede utilizar Search-R1 para entrenar un modelo que puede razonar y buscar para manejar una variedad de tareas.

escenario de aplicación

1. experimento de investigación
   Los investigadores pueden utilizar Search-R1 para reproducir los resultados del artículo y explorar la aplicación del aprendizaje por refuerzo en el entrenamiento de modelos.
2. Desarrollo de asistentes inteligentes
   Los desarrolladores pueden formar modelos que se integren en las herramientas de chat para ofrecer capacidades de búsqueda y razonamiento.
3. Consultas sobre conocimientos
   Los usuarios pueden utilizarla para responder rápidamente a preguntas complejas y obtener información actualizada mediante búsquedas.

CONTROL DE CALIDAD

1. ¿Qué modelos admite Search-R1?
   Actualmente se soportan los modelos base Qwen2.5-3B y Llama3.2-3B, otros modelos deben ser adaptados por ti mismo.
2. ¿Cuánto dura la formación?
   Dependiendo del conjunto de datos y del hardware, el conjunto de datos NQ tarda unas horas en entrenarse en una GPU gráfica de 24 GB.
3. ¿Cómo puedo comprobar la eficacia de mi formación?
   comprobar <Preliminary results> o consulta el registro de Wandb.