动手学习RAG:大模型重排模型 bge-reranker-v2-gemma微调

在这里插入图片描述

1. 环境准备

pip install transformers
pip install open-retrievals

2. 使用大模型做重排

from retrievals import LLMRanker

model_name = 'BAAI/bge-reranker-v2-gemma'
model = LLMRanker.from_pretrained(
            model_name,
            causal_lm=True,
            use_fp16=True,
        )

score = model.compute_score(['query', 'passage'])
print(score)

scores = model.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
print(scores)

请添加图片描述

3. 微调

MODEL_NAME='BAAI/bge-reranker-v2-gemma'
TRAIN_DATA="/root/kag101/src/open-retrievals/t2/t2_ranking.jsonl"
OUTPUT_DIR="/root/kag101/src/open-retrievals/t2/ft_out"

torchrun --nproc_per_node 1 \
    -m retrievals.pipelines.rerank \
    --output_dir ${OUTPUT_DIR} \
    --overwrite_output_dir \
    --model_name_or_path $MODEL_NAME \
    --model_type llm \
    --causal_lm True \
    --use_lora True \
    --data_name_or_path $TRAIN_DATA \
    --task_prompt "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'." \
    --query_instruction "A: " \
    --document_instruction 'B: ' \
    --positive_key positive \
    --negative_key negative \
    --learning_rate 2e-4 \
    --num_train_epochs 3 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 16 \
    --dataloader_drop_last True \
    --max_len 256 \
    --train_group_size 4 \
    --logging_steps 10 \
    --save_steps 20000 \
    --save_total_limit 1 \
    --bf16

请添加图片描述

4. 评测

在C-MTEB中进行评测。微调前保留10%的数据集作为测试集验证

微调前的指标:
请添加图片描述

微调后的指标:

请添加图片描述

{
  "dataset_revision": null,
  "mteb_dataset_name": "CustomReranking",
  "mteb_version": "1.1.1",
  "test": {
    "evaluation_time": 77.35,
    "map": 0.7057362287508586,
    "mrr": 0.8166538440773136
  }
}

微调后map从0.637上升至0.706,mrr从0.734上升至0.816

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

YueTann

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值