Search-R1: обучение с подкреплением для обучения больших моделей для поиска и рассуждений

Общее введение

Search-R1 - это проект с открытым исходным кодом, разработанный PeterGriffinJin на GitHub и построенный на фреймворке veRL. Он использует методы обучения с подкреплением (RL) для обучения больших языковых моделей (LLM), позволяя моделям автономно учиться рассуждать и вызывать поисковые системы для решения задач. Проект поддерживает такие базовые модели, как Qwen2.5-3B и Llama3.2-3B, а также расширяет DeepSeek-R1 и методы TinyZero. Пользователи могут использовать его для обучения моделей для решения однораундовых или многораундовых задач, при этом предоставляется код, наборы данных и журналы экспериментов. Официальный сайтобсудить статью или диссертацию (старая)Выпущенная в марте 2025 года модель проекта и данные доступны для скачивания на сайте Hugging Face для исследователей и разработчиков.

Список функций

• Обучение больших моделей с помощью обучения с подкреплением для улучшения рассуждений и поиска.
• Поддержка вызова API Google, Bing, Brave и других поисковых систем.
• Предоставляет возможности настройки 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 (поддерживает CUDA 12.1), введя следующую команду:

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

3. Установка vLLM
   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. Установка среды Retriever (опционально)
   Если требуется локальный сервер извлечения, создайте другую среду:

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Введите вопрос, который вы хотите задать.

предостережение

• Для обучения требуется графический процессор с видеопамятью не менее 24 ГБ (например, NVIDIA A100).
• Убедитесь, что ключ API действителен, а сетевое соединение стабильно.
• Официальные документы и лабораторные журналы (<Full experiment log 1> ответить пением <Full experiment log 2>) Предоставьте более подробную информацию.

Выполнив эти шаги, вы сможете использовать Search-R1 для обучения модели, способной рассуждать и искать, для решения различных задач.

сценарий применения

1. исследовательский эксперимент
   Исследователи могут воспроизвести результаты статьи с помощью Search-R1, чтобы изучить применение обучения с подкреплением в обучении моделей.
2. Разработка интеллектуального помощника
   Разработчики могут обучать модели для интеграции в инструменты чата, чтобы обеспечить возможности поиска и рассуждений.
3. Запросы знаний
   С его помощью пользователи могут быстро отвечать на сложные вопросы и получать актуальную информацию через поиск.

QA

1. Какие модели поддерживает Search-R1?
   В настоящее время поддерживаются базовые модели Qwen2.5-3B и Llama3.2-3B, другие модели необходимо адаптировать самостоятельно.
2. Как долго длится обучение?
   В зависимости от набора данных и аппаратного обеспечения, обучение на наборе данных NQ занимает около нескольких часов на графическом процессоре с 24 ГБ памяти.
3. Как проверить эффективность обучения?
   проверять <Preliminary results> или проверьте журнал Wandb.