0. 克隆仓库。
1. 根据您的 CUDA 版本，从 PyTorch 安装 PyTorch。
2. 安装依赖项和 Tevatron。

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

0. 克隆仓库。
1. 按照 官方指南 安装 JAX。
2. 安装依赖项。

pip install transformers datasets
pip install flax optax

3. 安装 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 ..

4. 安装 Tevatron。

pip install -e .

要在 GPU 上运行 Tevatron 的 JAX 实现，我们建议使用 NVIDIA 提供的 jax-toolbox jax container 镜像。

以下是基于 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 接受 jsonl 格式的训练或推理数据，每行是一个 JSON 对象，格式如下：

1. 训练数据

{
   "query_id": "<query id>",
   "query_text": "<query text>",
   "query_image": "<query image>",
   "positive_document_ids": ["<passage id>", ...],
   "negative_document_ids": ["<passage id>", ...],
}

其中，positive_passages 是与查询相关联的标注相关段落，而 negative_passages 通常是来自检索系统（如 BM25、DPR）顶部结果中的不相关（硬负样本）段落。对于 QA 数据集，还可以包含 answers 等其他字段。

2. 语料数据

{
   "docid": "<document id>",
   "document_text": "<document text>",
   "document_image": "<document image>",
}

其中，每行代表语料库中的一个文档。

需要注意的是，训练和语料数据中的图像字段是可选的，可以省略（即纯文本模态的检索）。

自包含数据集

Tevatron 自带了几种常用的神经检索数据集。 （通过 HuggingFace 提供）。 这些数据集可以在训练和编码过程中通过设置 --dataset_name <hgf dataset name> 自动下载。

在本示例中，我们将使用自包含数据集 Tevatron/msmarco-passage-aug 进行训练，其硬负样本是从 BM25 前 200 名和 CoCondenser 前 200 名结果中混合采样的。

训练

deepspeed --include localhost:0,1,2,3 --master_port 60000 --module tevatron.retriever.driver.train \
  --deepspeed deepspeed/ds_zero3_config.json \
  --output_dir retriever-mistral \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --lora \
  --lora_target_modules q_proj,k_proj,v_proj,o_proj,down_proj,up_proj,gate_proj \
  --save_steps 50 \
  --dataset_name Tevatron/msmarco-passage-aug \
  --query_prefix "Query: " \
  --passage_prefix "Passage: " \
  --bf16 \
  --pooling eos \
  --append_eos_token \
  --normalize \
  --temperature 0.01 \
  --per_device_train_batch_size 8 \
  --gradient_checkpointing \
  --train_group_size 16 \
  --learning_rate 1e-4 \
  --query_max_len 32 \
  --passage_max_len 156 \
  --num_train_epochs 1 \
  --logging_steps 10 \
  --overwrite_output_dir \
  --gradient_accumulation_steps 4

每次查询处理的段落数：8×4×16 = 512

每次更新的查询数：8×4×4 = 128

上述训练配置在 4 张 A6000 GPU 上大约需要 70 小时。

如果仅使用 1 张 A100 GPU，则相当于约 110 小时的训练时间。

编码

查询编码

EMBEDDING_OUTPUT_DIR=<用于保存查询嵌入的文件夹>
CUDA_VISIBLE_DEVICES=4 python -m tevatron.retriever.driver.encode \
  --output_dir=temp \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --lora_name_or_path retriever-mistral \
  --lora \
  --query_prefix "Query: " \
  --passage_prefix "Passage: " \
  --bf16 \
  --pooling eos \
  --append_eos_token \
  --normalize \
  --encode_is_query \
  --per_device_eval_batch_size 128 \
  --query_max_len 32 \
  --passage_max_len 156 \
  --dataset_name Tevatron/msmarco-passage \
  --dataset_split dev \
  --encode_output_path $EMBEDDING_OUTPUT_DIR/query-dev.pkl

语料库编码

EMBEDDING_OUTPUT_DIR=<用于保存查询嵌入的文件夹>
for s in 0 1 2 3
do
gpuid=$s
CUDA_VISIBLE_DEVICES=$gpuid python -m tevatron.retriever.driver.encode \
  --output_dir=temp \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --lora_name_or_path retriever-mistral \
  --lora \
  --query_prefix "Query: " \
  --passage_prefix "Passage: " \
  --bf16 \
  --pooling eos \
  --append_eos_token \
  --normalize \
  --per_device_eval_batch_size 128 \
  --query_max_len 32 \
  --passage_max_len 156 \
  --dataset_name Tevatron/msmarco-passage-corpus \
  --dataset_number_of_shards 4 \
  --dataset_shard_index ${s} \
  --encode_output_path $EMBEDDING_OUTPUT_DIR/corpus.${s}.pkl
done

在命令末尾添加 & 可以在后台并行运行。

检索

set -f && python -m tevatron.retriever.driver.search \
    --query_reps $EMBEDDING_OUTPUT_DIR/query-dev.pkl \
    --passage_reps $EMBEDDING_OUTPUT_DIR/corpus*.pkl \
    --depth 1000 \
    --batch_size 64 \
    --save_text \
    --save_ranking_to $EMBEDDING_OUTPUT_DIR/run.dev.txt

输出文件的格式为每行 <query_id> <passage_id> <score>。

训练

对于 GPU 训练，设置 XLA_PYTHON_CLIENT_MEM_FRACTION=.95，并且如果已安装 TransformersEngine，则确保查询和段落长度是 64 的倍数。

python -m tevatron.tevax.experimental.mp.train_lora  \
   --checkpoint_dir retriever-mistral-jax \
   --train_file Tevatron/msmarco-passage-aug \
   --model_name mistralai/Mistral-7B-v0.1 \
   --model_type mistral \
   --batch_size 128 \
   --num_target_passages 16 \
   --learning_rate 1e-4 \
   --seed 12345 \
   --mesh_shape 1 -1 \
   --weight_decay 0.00001 \
   --num_epochs 1 \
   --max_query_length 64 \
   --max_passage_length 128 \
   --pooling eos \
   --scale_by_dim True \
   --grad_cache \
   --passage_num_chunks 32 \
   --query_num_chunks 4

每个查询的批次数：128x16 = 2048

每次更新的查询数量：128

上述训练设置在 v4-8 TPU VM 上大约需要 35 小时。

在 1xA100 GPU 上进行等效训练大约需要 80 小时。

编码

查询编码

python -m tevatron.tevax.experimental.mp.encode  \
   --model_type mistral \
   --model_name_or_path mistralai/Mistral-7B-v0.1 \
   --model_config_name_or_path mistralai/Mistral-7B-v0.1 \
   --tokenizer_name_or_path mistralai/Mistral-7B-v0.1 \
   --dataset_name_or_path Tevatron/msmarco-passage \
   --split dev \
   --output_dir $EMBEDDING_OUTPUT_DIR/query-embedding \
   --batch_size 32 \
   --input_type query \
   --max_seq_length 64 \
   --mesh_shape 1 -1 \
   --lora retriever-mistral-jax/lora \
   --scale_by_dim

语料库编码

python -m tevatron.tevax.experimental.mp.encode  \
   --model_type mistral \
   --model_name_or_path mistralai/Mistral-7B-v0.1 \
   --model_config_name_or_path mistralai/Mistral-7B-v0.1 \
   --tokenizer_name_or_path mistralai/Mistral-7B-v0.1 \
   --dataset_name_or_path Tevatron/msmarco-passage-corpus \
   --output_dir $EMBEDDING_OUTPUT_DIR/corpus-embedding \
   --batch_size 32 \
   --input_type passage \
   --max_seq_length 128 \
   --mesh_shape 1 -1 \
   --lora retriever-mistral-jax/lora \
   --scale_by_dim

检索

set -f && python -m tevatron.retriever.driver.search \
    --query_reps $EMBEDDING_OUTPUT_DIR/query-embedding/*.pkl \
    --passage_reps $EMBEDDING_OUTPUT_DIR/corpus-embedding/*.pkl \
    --depth 1000 \
    --batch_size 64 \
    --save_text \
    --save_ranking_to $EMBEDDING_OUTPUT_DIR/run.dev.txt

输出文件的格式为每行 <query_id> <passage_id> <score>。