• 動手學習RAG: 向量模型
• 動手學習RAG: BGE向量模型微調實踐]()
• 動手學習RAG: BCEmbedding 向量模型 微調實踐]()
• BCE ranking 微調實踐]()
• GTE向量與排序模型 微調實踐]()
• 模型微調中的模型序列長度]()
• 相似度與溫度系數(shù)

本文我們來進行ColBERT模型的實踐，按慣例，還是以open-retrievals中的代碼為藍本。在RAG興起之后，ColBERT也獲得了更多的關注。ColBERT整體結構和雙塔特別相似，但遲交互式也就意味著比起一般ranking模型，交互來的更晚一些。

準備環(huán)境

pip install transformers
pip install open-retrievals

準備數(shù)據(jù)

還是采用C-MTEB/T2Reranking數(shù)據(jù)。

• 每個樣本有query, positive, negative。其中query和positive構成正樣本對，query和negative構成負樣本對
  

使用

由于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

• 生成向量的方法

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)

• 計算句子對 相似度的方法

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)

微調

嘗試了兩種方法來做，一種是調包自己寫代碼，一種是采用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模型權重

損失函數(shù)下降

{'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': 471.8191, 'train_samples_per_second': 33.827, 'train_steps_per_second': 1.055, 'train_loss': 2.4146631079984, 'epoch': 3.0}

評測

在C-MTEB中進行評測。微調前保留10%的數(shù)據(jù)集作為測試集驗證

from datasets import load_datasetdataset = 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)

微調前的指標：

微調后的指標：

{"dataset_revision": null,"mteb_dataset_name": "CustomReranking","mteb_version": "1.1.1","test": {"evaluation_time": 221.45,"map": 0.6950128151840831,"mrr": 0.8193114944390455}
}

番外：從語言模型直接訓練ColBERT

之前的例子里是從BAAI/bge-m3繼續(xù)微調，這里再跑一個從hfl/chinese-roberta-wwm-ext訓練一個ColBERT模型

• 注意，從頭跑需要設置更大的學習率與更多的epochs

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/srctorchrun --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

微調后指標

{"dataset_revision": null,"mteb_dataset_name": "CustomReranking","mteb_version": "1.1.1","test": {"evaluation_time": 75.38,"map": 0.6865308507184888,"mrr": 0.8039965986394558}
}