Search-R1: Aprendizado por reforço para treinar modelos grandes para pesquisa e raciocínio

Introdução geral

O Search-R1 é um projeto de código aberto desenvolvido por PeterGriffinJin no GitHub e baseado na estrutura veRL. Ele usa técnicas de aprendizagem por reforço (RL) para treinar grandes modelos de linguagem (LLMs), permitindo que os modelos aprendam de forma autônoma a raciocinar e invocar mecanismos de pesquisa para resolver problemas. O projeto oferece suporte a modelos básicos, como o Qwen2.5-3B e o Llama3.2-3B, estendendo o DeepSeek-R1 e TinyZero. Os usuários podem usá-lo para treinar modelos para lidar com tarefas de uma ou várias rodadas, com código, conjuntos de dados e registros de experimentos fornecidos. Oficialdiscutir um artigo ou tese (antigo)Lançados em março de 2025, o modelo e os dados do projeto estão disponíveis para download no site Hugging Face para pesquisadores e desenvolvedores.

Lista de funções

• Treinamento de modelos grandes por meio do aprendizado por reforço para aprimorar o raciocínio e a pesquisa.
• Suporte para chamar a API do Google, Bing, Brave e outros mecanismos de pesquisa.
• Oferece recursos de ajuste de LoRA e ajuste fino supervisionado para otimizar o desempenho do modelo.
• Reordenador pronto para uso incorporado para melhorar a precisão dos resultados de pesquisa.
• Inclui registros de laboratório detalhados e documentos para apoiar a reprodução dos resultados.
• Oferece funcionalidade de servidor de pesquisa local para facilitar a personalização das pesquisas.
• Suporte para upload de conjuntos de dados e corpora personalizados pelo usuário.

Usando a Ajuda

O Search-R1 é destinado a usuários com conhecimentos básicos de programação e aprendizado de máquina. Abaixo está um guia detalhado de instalação e uso para que você possa começar rapidamente.

Processo de instalação

Para usar o Search-R1, você precisa primeiro configurar o ambiente. As etapas são as seguintes:

1. Criação de um ambiente Search-R1
   É executado no terminal:

conda create -n searchr1 python=3.9
conda activate searchr1

Isso cria um ambiente virtual Python 3.9.

2. Instalando o PyTorch
   Instale o PyTorch 2.4.0 (compatível com CUDA 12.1) digitando o seguinte comando:

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

3. Instalação do vLLM
   vLLM é uma biblioteca essencial para a execução de modelos grandes, instale a versão 0.6.3:

pip3 install vllm==0.6.3

As versões 0.5.4, 0.4.2 ou 0.3.1 também estão disponíveis.

4. Instalação do veRL
   Execute-o no diretório raiz do projeto:

pip install -e .

Isso instalará a estrutura veRL.

5. Instalar dependências opcionais
   Para melhorar o desempenho, instale o Flash Attention e o Wandb:

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

6. Instalação do Retriever Environment (opcional)
   Se for necessário um servidor de recuperação local, crie outro ambiente:

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

Início rápido

A seguir estão as etapas para treinar um modelo com base no conjunto de dados NQ:

1. Download de índices e corpora
   Defina o caminho para salvar e execute-o:

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. Processamento de dados NQ
   Execute o script para gerar dados de treinamento:

python scripts/data_process/nq_search.py

3. Iniciar o servidor de recuperação
   Funciona no ambiente do Retriever:

conda activate retriever
bash retrieval_launch.sh

4. Treinamento de corrida RL
   Execute em um ambiente Search-R1:

conda activate searchr1
bash train_ppo.sh

Isso usará o modelo básico Llama-3.2-3B para treinamento de PPO.

Uso de conjuntos de dados personalizados

1. Preparação de dados de controle de qualidade
   Os dados precisam estar no formato JSONL e cada linha deve conter os seguintes 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. Preparando o corpus
   O corpus precisa estar no formato JSONL, com cada linha contendo id responder cantando contentsComo:

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

   referenciável <example/corpus.jsonl>.

3. Corpus de indexação (opcional)
   Se estiver usando a pesquisa local, execute:

   bash search_r1/search/build_index.sh
   

Chamada de um mecanismo de pesquisa personalizado

1. modificações <search_r1/search/retriever_server.py>As APIs de configuração.
2. Inicie o servidor:

   python search_r1/search/retriever_server.py
   

3. O modelo é passado pelo http://127.0.0.1:8000/retrieve Pesquisa de chamadas.

operação de inferência

1. Inicie o servidor de recuperação:

   bash retrieval_launch.sh
   

2. Raciocínio em execução:

   python infer.py
   

3. modificações <infer.py> Linha 7 questiondigite a pergunta que você deseja fazer.

advertência

• O treinamento requer uma GPU com pelo menos 24 GB de memória de vídeo (por exemplo, NVIDIA A100).
• Verifique se a chave da API é válida e se a conexão de rede é estável.
• Documentos oficiais e revistas de laboratório (<Full experiment log 1> responder cantando <Full experiment log 2>) Forneça mais detalhes.

Com essas etapas, você pode usar o Search-R1 para treinar um modelo que possa raciocinar e pesquisar para lidar com uma variedade de tarefas.

cenário do aplicativo

1. experimento de pesquisa
   Os pesquisadores podem usar o Search-R1 para reproduzir os resultados do artigo e explorar a aplicação do aprendizado por reforço no treinamento de modelos.
2. Desenvolvimento de assistentes inteligentes
   Os desenvolvedores podem treinar modelos a serem integrados em ferramentas de bate-papo para fornecer recursos de pesquisa e raciocínio.
3. Consultas de conhecimento
   Os usuários podem usá-lo para responder rapidamente a perguntas complexas e obter informações atualizadas por meio de pesquisa.

QA

1. Quais modelos são compatíveis com o Search-R1?
   Atualmente, há suporte para os modelos básicos Qwen2.5-3B e Llama3.2-3B; outros modelos precisam ser adaptados por você.
2. Quanto tempo dura o treinamento?
   Dependendo do conjunto de dados e do hardware, o conjunto de dados NQ leva cerca de algumas horas para ser treinado em uma GPU gráfica de 24 GB.
3. Como posso verificar a eficácia do meu treinamento?
   confira <Preliminary results> ou verifique o registro do Wandb.