使用RAG技术构建企业级文档问答系统：检索优化: Rerank模型微调

最新推荐文章于 2025-04-25 15:59:36 发布

大靠山

最新推荐文章于 2025-04-25 15:59:36 发布

阅读量2.6k

点赞数 21

文章标签：人工智能学习分类大数据 python

本文链接：https://blog.csdn.net/m0_59235245/article/details/143131241

版权

1 概述

关于Reranker的作用，它能够使得检索到的上下文有更准确的排序，因而常常可以大幅提升RAG的最终效果。

https://cohere.com/blog/rerank

对Reranker模型进行二次训练，也是分为二次预训练和微调，与Embedding模型类似，二次预训练通常情况下不会带来大幅度的性能提升。本文以常用的BGE系列reranker为例。

下图最后一行是增加bge-reranker-base微调后的的结果，从检索效果上看，微调后的模型，性能已经接近比它大许多倍的bge-reranker-v2-m3了。

在这里插入图片描述

2 环境准备

Reranker模型的训练，虽然对机器的性能要求没有对训练LLM那么高，但也还是有一定要求的，GPU是需要的，下面是本文所使用的软硬件环境，供大家参考。

2.1 硬件环境

本文所用硬件环境如下：

CPU：i7-9700K
内存：64GB
GPU：GTX 1080Ti（11G显存）

2.2 软件环境

软件环境列举主要的Python依赖：

Python：3.10.9
pytorch包：2.2.1
FlagEmbedding包：1.2.10

3 核心代码

本文对应完整代码地址：https://github.com/Steven-Luo/MasteringRAG/blob/main/retrieval/06_reranker_ft.ipynb

3.1 微调数据构建

此处数据构建方法。

def build\_qa\_samples(df, neg\_batch\_size\=-1, n\_neg\_batch\=5):  
    """  
    构建qa样本  
    :param df: 包含qa的DataFrame，共两列，question和answer  
    :param neg\_batch\_size: 负样本数量，为-1时表示将所有负样本和单个正样本配对，否则会将负样本拆开，结果中的query可能会重复  
    """  
    from tqdm.auto import tqdm  
    import math  
  
    data = \[\]  
    for idx, row in tqdm(df.iterrows(), total\=len(df)):  
        question = row\['question'\]  
        answer = row\['answer'\]  
        \# 筛选同category的，增加难度  
        neg\_samples = df\[df\['question'\] != question\]\['answer'\].values.tolist()  
        neg\_batch\_count = math.ceil((len(df) - 1) / neg\_batch\_size)  
        neg\_batch\_count = min(n\_neg\_batch, neg\_batch\_count)  
        for neg\_batch\_idx in range(neg\_batch\_count):  
            batch\_neg\_samples = neg\_samples\[neg\_batch\_idx \* neg\_batch\_size: (neg\_batch\_idx + 1) \* neg\_batch\_size\]  
            batch\_neg\_samples = \[item for item in batch\_neg\_samples if item != answer\]  
            data.append({  
                'query': question,  
                'pos': \[answer\],  
                'neg': batch\_neg\_samples  
            })  
    return data  
  
def write\_samples(samples, save\_filename):  
    import json  
  
    with open(save\_filename, 'w') as f:  
        for sample in samples:  
            f.write(json.dumps(sample, ensure\_ascii\=False))  
            f.write('\\n')

build_qa_samples函数是个通用模板，通常情况下，Q-D的微调效果是更好的，此处也是一样，在调用build_qa_samples函数前，使用下面的语句将context赋值给answer，这样构建的样本就是Q-D的样本对了：

df\['answer'\] = df\['context'\]  
qd\_samples = build\_qa\_samples(qa\_df, neg\_batch\_size=16, n\_neg\_batch=32)

其中df包含原始的问题、上下文、参考答案，构造方法参考使用RAG技术构建企业级文档问答系统之QA抽取，构造好的数据样例如下：

在这里插入图片描述

build_qa_samples函数的返回结果样例如下：

{"query": "美元指数在2023年的走势如何？",  
 "pos": \["美元指数高位震荡后走弱"\],  
 "neg": \["全球货物贸易量指数和价格指数下行，主要经济体出口贸易同比增速下降。",  
  "欧美央行货币政策延续收紧态势，但步伐整体放缓。",  
  "高利率环境抑制债券融资需求，债券违约风险持续上升，美国政府债务可持续性问题引发市场关注。",  
  "展望2024年，预计全球经济复苏将依旧疲软，主要经济体增长态势和货币政策将进一步分化。",  
  "发达经济体增速明显放缓，预计2023年增速较2022年下降1个百分点。",  
  "新兴经济体增速与2022年大致持平，预计2023年增速比2022年下降0.1个百分点。",  
  "美国GDP环比增长折年率为4.9%，比二季度增速高2.8个百分点。",  
  "其中，主要新兴经济体工业生产指数普遍走高，如俄罗斯、土耳其、南非等，而发达经济体中的美国和韩国回升，英国、德国、意大利下行，日本波动较大，整体趋于平稳。",  
  "主要新兴经济体工业生产指数普遍走高，如俄罗斯、土耳其、南非等",  
  "全球融资环境收紧和经济下行压力对工业生产前景带来较大影响",  
  "欧洲各国消费指数整体维持稳定（图4）",  
  "是上半年免于陷入衰退的主要动力。",  
  "OECD消费者信心指数从7月开始连续3个月回落",  
  "美国私人投资在2023年一季度触底后逐渐反弹",  
  "2023年二季度，欧元区固定资本形成总额环比增长0.1%，比一季度增速下降0.3个百分点",  
  "房地产对GDP环比增长拉动率转为负值"\]}

3.2 bge-reranker-base微调脚本

将下方脚本保存为finetune_bge_reranker_base_v1.sh，然后在shell中执行即可：

#!/bin/bash  
  
SCRIP\_DIR\=$(echo \`cd $(dirname $0); pwd\`)  
  
export TRAIN\_DATASET\=outputs/v1\_20240713/emb\_samples\_qd\_v2.jsonl  
  
export N\_EPOCH\=2  
export TRAIN\_GROUP\_SIZE\=32  
  
export GRADIENT\_ACCUMULATION\_STEPS\=32  
export PER\_DEVICE\_TRAIN\_BATCH\_SIZE\=1  
export N\_NODES\=1  
export BATCH\_SIZE\=\`expr ${GRADIENT\_ACCUMULATION\_STEPS} \\\* ${PER\_DEVICE\_TRAIN\_BATCH\_SIZE} \\\* ${N\_NODES}\`  
  
export VERSION\=ft\_v1\_bge\_base\_epoch\_${N\_EPOCH}\_bz\_${BATCH\_SIZE}\_trgrp\_${TRAIN\_GROUP\_SIZE}\_$(date +"%Y%m%d\_%H%M")  
  
\# 如果需要将训练进度上传Wandb，可以取消下方注释  
#export WANDB\_PROJECT=RAG-From-Scratch-Reranker-Finetune  
#export WANDB\_API\_KEY=替换为自己的Key  
#export WANDB\_NAME=${VERSION}  
  
export OUTPUT\_DIR\=experiments/reranker/finetune/${VERSION}  
  
if \[ ! -d "${OUTPUT\_DIR}" \]; then  
    mkdir -p "${OUTPUT\_DIR}"  
fi  
  
torchrun --nproc\_per\_node ${N\_NODES} \\  
\-m FlagEmbedding.reranker.run \\  
\--output\_dir ${OUTPUT\_DIR} \\  
\--model\_name\_or\_path BAAI/bge-reranker-base \\  
\--train\_data ${TRAIN\_DATASET} \\  
\--learning\_rate 1e-5 \\  
\--fp16 \\  
\--num\_train\_epochs ${N\_EPOCH} \\  
\--per\_device\_train\_batch\_size ${PER\_DEVICE\_TRAIN\_BATCH\_SIZE} \\  
\--gradient\_accumulation\_steps ${GRADIENT\_ACCUMULATION\_STEPS} \\  
\--dataloader\_drop\_last True \\  
\--train\_group\_size ${TRAIN\_GROUP\_SIZE} \\  
\--logging\_steps 5 \\  
\--save\_steps 50 \\  
\--save\_total\_limit 10 \\  
\--warmup\_ratio 0.05 \\  
\--lr\_scheduler\_type cosine  
  
  
cp "$SCRIP\_DIR/$0" ${OUTPUT\_DIR}

4 训练过程

训练过程的完整监控，可以访问下方链接查看：

https://api.wandb.ai/links/steven-luog/og13o7yp

以下是截图：

在这里插入图片描述

5 模型使用

模型已经上传至HuggingFace，模型ID为stevenluo/bge-reranker-base-ft-v1，像正常BGE Reranker模型一样使用即可，示例代码如下：

from FlagEmbedding import FlagModel  
sentences\_1 = \["样例数据-1", "样例数据-2"\]  
sentences\_2 = \["样例数据-3", "样例数据-4"\]  
model = FlagModel(  
    'stevenluo/bge-reranker-base-ft-v1',   
    query\_instruction\_for\_retrieval="为这个句子生成表示以用于检索相关文章：",  
    use\_fp16=True)  
embeddings\_1 = model.encode(sentences\_1)  
embeddings\_2 = model.encode(sentences\_2)  
similarity = embeddings\_1 @ embeddings\_2.T  
print(similarity)

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述