tevatron - 神经检索模型训练与推理工具

Tevatron V2 项目介绍

Tevatron 项目的目标是提供一个灵活高效的工具包，能够在大规模的环境下进行神经网络模型的训练和推理，特别是用于信息检索的模型。

功能特点

Tevatron V2 具备以下主要功能特点：

大规模神经检索器训练：支持在 GPU 和 TPU 上对十亿规模的语言模型（LLM）神经检索器进行训练。
LoRA 参数高效调优：支持使用 LoRA 进行参数高效的微调。
集成多种高效训练技术：包括与 DeepSpeed、闪存注意力机制、梯度累积等高效训练技术的集成。
内置数据集：适用于神经检索和开放域问答任务的自包含数据集。
直接加载和微调最先进的预训练模型：例如 BGE-Embedding 和 Instruct-E5，可以从 HuggingFace 平台直接加载。

安装指南

根据所用的硬件平台，Tevatron 提供多个安装选项：

PyTorch (GPU)

克隆代码库。
按照相应的 CUDA 版本从 PyTorch 网站安装 PyTorch。

安装依赖项和 Tevatron。

pip install transformers datasets peft
pip install deepspeed accelerate
pip install faiss
pip install -e .

JAX (TPU)

克隆代码库。
按照官方指南安装 JAX。

安装所需的依赖包。

pip install transformers datasets
pip install flax optax

安装 Magix 和 GradCache。

git clone https://github.com/luyug/magix.git
cd magix && pip install -e . && cd ..
git clone https://github.com/luyug/GradCache.git
cd GradCache && pip install -e . && cd ..

安装 Tevatron。
```
pip install -e .
```

JAX (GPU)

使用 NVIDIA 提供的 jax-toolbox 容器镜像运行 JAX 实现的 Tevatron。

示例 Dockerfile：

FROM ghcr.io/nvidia/jax:jax-2024-03-08

RUN apt-get update && \
    apt-get install -y --no-install-recommends python3-pip && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    pip install --no-cache-dir transformers sentencepiece simple_parsing datasets orbax==0.4.8 && \
    pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu

RUN git clone https://github.com/luyug/magix.git && \
    cd magix && pip install -e . && cd .. && \
    git clone https://github.com/luyug/GradCache.git \
    cd GradCache && pip install -e . && cd .. \
    git clone https://github.com/texttron/tevatron.git && \
    cd tevatron && pip install -e .

Tevatron 101

通过一个简单的例子展示如何使用 Tevatron 对 Mistral-7B 模型进行 LoRA 微调并在 MS MARCO 数据集上训练得到的检索器。

数据准备

Tevatron 的输入数据采用 jsonl 格式，每行是一个 JSON 对象，例如训练数据：

{
   "query_id": "<query id>",
   "query": "<query text>",
   "positive_passages": [
     {"docid": "<passage id>", "title": "<passage title>", "text": "<passage body>"}
   ],
   "negative_passages": [
     {"docid": "<passage id>", "title": "<passage title>", "text": "<passage body>"}
   ]
}

自包含数据集支持神经检索，将使用 Tevatron/msmarco-passage-aug 数据集进行训练。

训练与编码

详细的训练和编码步骤分别适用于 PyTorch 和 JAX 环境，具体命令和参数设置详见安装部分。

引用与致谢

如果您觉得 Tevatron 有帮助，请参考以下论文：

@article{Gao2022TevatronAE,
  title={Tevatron: An Efficient and Flexible Toolkit for Dense Retrieval},
  ...
}

项目感谢所有依赖库的贡献者，同时感谢 Google 的 TPU 研究云提供的支持。

联系方式

如对工具包有疑问，可以在平台上提出 issue，或者通过以下邮箱联系开发团队：

Luyu Gao: luyug@cs.cmu.edu
Xueguang Ma: x93ma@uwaterloo.ca

希望这篇介绍能帮助您更好地了解和使用 Tevatron 项目。