- 动手学习RAG: 向量模型
- 动手学习RAG: BGE向量模型微调实践]()
- 动手学习RAG: BCEmbedding 向量模型 微调实践]()
- BCE ranking 微调实践]()
- GTE向量与排序模型 微调实践]()
- 模型微调中的模型序列长度]()
本文我们来进行ColBERT模型的实践，按惯例，还是以open-retrievals中的代码为蓝本。在RAG兴起之后，ColBERT也获得了更多的关注。ColBERT整体结构和双塔特别相似，但迟交互式也就意味着比起一般ranking模型，交互来的更晚一些。
Screen Shot 2024-09-12 at 19.07.51.png
本文代码：https://colab.research.google.com/drive/1QVtqhQ080ZMltXoJyODMmvEQYI6oo5kO?usp=sharing
准备环境
pip install transformers
pip install open-retrievals
准备数据
还是采用C-MTEB/T2Reranking数据。
- 每个样本有query, positive, negative。其中query和positive构成正样本对，query和negative构成负样本对
[图片上传失败...(image-583d0-1726122973759)]
使用
由于ColBERT作为迟交互式模型，既可以像向量模型一样生成向量，也可以计算相似度。BAAI/bge-m3中的colbert模型是基于XLMRoberta训练而来，因此使用ColBERT可以直接从bge-m3中加载预训练权重。
import transformers
from retrievals import ColBERT
model_name_or_path: str = 'BAAI/bge-m3'
model = ColBERT.from_pretrained(
model_name_or_path,
colbert_dim=1024,
use_fp16=True,
loss_fn=ColbertLoss(use_inbatch_negative=True),
)
model
[图片上传失败...(image-54e0cd-1726122973759)]
- 生成向量的方法
sentences_1 = ["In 1974, I won the championship in Southeast Asia in my first kickboxing match", "In 1982, I defeated the heavy hitter Ryu Long."]
sentences_2 = ['A dog is chasing car.', 'A man is playing a guitar.']
output_1 = model.encode(sentences_1, normalize_embeddings=True)
print(output_1.shape, output_1)
output_2 = model.encode(sentences_2, normalize_embeddings=True)
print(output_2.shape, output_2)
[图片上传失败...(image-18c91a-1726122973759)]
- 计算句子对 相似度的方法
sentences = [
["In 1974, I won the championship in Southeast Asia in my first kickboxing match", "In 1982, I defeated the heavy hitter Ryu Long."],
["In 1974, I won the championship in Southeast Asia in my first kickboxing match", 'A man is playing a guitar.'],
]
scores_list = model.compute_score(sentences)
print(scores_list)
[图片上传失败...(image-d2a27b-1726122973759)]
微调
尝试了两种方法来做，一种是调包自己写代码，一种是采用open-retrievals中的代码写shell脚本。这里我们采用第一种，另外一种方法可参考文章最后番外中的微调
import transformers
from transformers import AutoTokenizer, TrainingArguments, get_cosine_schedule_with_warmup, AdamW
from retrievals import AutoModelForRanking, RerankCollator, RerankTrainDataset, RerankTrainer, ColBERT, RetrievalTrainDataset, ColBertCollator
from retrievals.losses import ColbertLoss
transformers.logging.set_verbosity_error()
model_name_or_path: str = 'BAAI/bge-m3'
learning_rate: float = 1e-5
batch_size: int = 2
epochs: int = 1
output_dir: str = './checkpoints'
train_dataset = RetrievalTrainDataset(
'C-MTEB/T2Reranking', positive_key='positive', negative_key='negative', dataset_split='dev'
)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
data_collator = ColBertCollator(
tokenizer,
query_max_length=64,
document_max_length=128,
positive_key='positive',
negative_key='negative',
)
model = ColBERT.from_pretrained(
model_name_or_path,
colbert_dim=1024,
loss_fn=ColbertLoss(use_inbatch_negative=False),
)
optimizer = AdamW(model.parameters(), lr=learning_rate)
num_train_steps = int(len(train_dataset) / batch_size * epochs)
scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=0.05 * num_train_steps, num_training_steps=num_train_steps)
training_args = TrainingArguments(
learning_rate=learning_rate,
per_device_train_batch_size=batch_size,
num_train_epochs=epochs,
output_dir = './checkpoints',
remove_unused_columns=False,
gradient_accumulation_steps=8,
logging_steps=100,
)
trainer = RerankTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
data_collator=data_collator,
)
trainer.optimizer = optimizer
trainer.scheduler = scheduler
trainer.train()
model.save_pretrained(output_dir)
训练过程中会加载BAAI/bge-m3模型权重
[图片上传失败...(image-6ac845-1726122973759)]
损失函数下降
{'loss': 7.4858, 'grad_norm': 30.484981536865234, 'learning_rate': 4.076305220883534e-06, 'epoch': 0.6024096385542169}
{'loss': 1.18, 'grad_norm': 28.68316650390625, 'learning_rate': 3.072289156626506e-06, 'epoch': 1.2048192771084336}
{'loss': 1.1399, 'grad_norm': 14.203865051269531, 'learning_rate': 2.068273092369478e-06, 'epoch': 1.8072289156626506}
{'loss': 1.1261, 'grad_norm': 24.30337905883789, 'learning_rate': 1.0642570281124499e-06, 'epoch': 2.4096385542168672}
{'train_runtime': 465.7768, 'train_samples_per_second': 34.265, 'train_steps_per_second': 1.069, 'train_loss': 2.4146631079984, 'epoch': 3.0}
评测
在C-MTEB中进行评测。微调前保留10%的数据集作为测试集验证
from datasets import load_dataset
dataset = load_dataset("C-MTEB/T2Reranking", split="dev")
ds = dataset.train_test_split(test_size=0.1, seed=42)
ds_train = ds["train"].filter(
lambda x: len(x["positive"]) > 0 and len(x["negative"]) > 0
)
ds_train.to_json("t2_ranking.jsonl", force_ascii=False)
微调前的指标：
[图片上传失败...(image-5307ee-1726122973759)]
微调后的指标：
09/12/2024 15:30:26 - INFO - mteb.evaluation.MTEB - Evaluation for CustomReranking on test took 221.45 seconds
09/12/2024 15:30:26 - INFO - mteb.evaluation.MTEB - Scores: {'map': 0.6950128151840831, 'mrr': 0.8193114944390455, 'evaluation_time': 221.45}
番外：从语言模型直接训练ColBERT
之前的例子里是从BAAI/bge-m3继续微调，这里再跑一个从hfl/chinese-roberta-wwm-ext训练一个ColBERT模型
MODEL_NAME='hfl/chinese-roberta-wwm-ext'
TRAIN_DATA="/root/kaggle101/src/open-retrievals/t2/t2_ranking.jsonl"
OUTPUT_DIR="/root/kaggle101/src/open-retrievals/t2/ft_out"
cd /root/open-retrievals/src
torchrun --nproc_per_node 1 \
--module retrievals.pipelines.rerank \
--output_dir $OUTPUT_DIR \
--overwrite_output_dir \
--model_name_or_path $MODEL_NAME \
--tokenizer_name $MODEL_NAME \
--model_type colbert \
--do_train \
--data_name_or_path $TRAIN_DATA \
--positive_key positive \
--negative_key negative \
--learning_rate 5e-5 \
--bf16 \
--num_train_epochs 5 \
--per_device_train_batch_size 32 \
--dataloader_drop_last True \
--query_max_length 128 \
--max_length 256 \
--train_group_size 4 \
--unfold_each_positive false \
--save_total_limit 1 \
--logging_steps 100 \
--use_inbatch_negative False
网友评论