LLM大模型: RAG两大核心利器 — embedding和reranker模型微调fine-tune

程序学到昏

于 2024-08-06 17:50:41 发布

阅读量640

点赞数 9

文章标签： embedding 产品经理 langchain RAG 大模型 ai大模型微调

本文链接：https://blog.csdn.net/Y525698136/article/details/140962174

版权

要想RAG好，embedding和reranker必须给力！目前市面上流行的embedding和reranker使用的都是通用预料训练，并未针对安全这个细分领域定制，所以有必要使用安全领域的预料微调一下！目前所有的预料中，获取成本比较低、并且不需要专门投入人力标注的有两种：

网上各种安全论坛的博客、各大热门产品的漏洞说明等
用户的点赞反馈数据（chatGPT、copilot等都有该功能）

对于作者本人而言，用户的点赞反馈数据更容易获取，所以这里使用这类数据，借鉴RLHF-DPO的思路对embedding和reranker模型做微调！训练样本的数据格式如下：

所有资料 ⚡️ ，朋友们如果有需要全套《LLM大模型入门+进阶学习资源包》，扫码获取~

👉CSDN大礼包🎁：全网最全《LLM大模型入门+进阶学习资源包》免费分享（安全链接，放心点击）👈

{
  "query": "如何使用IDA Pro反汇编一个二进制文件？",
  "positive": [
    "使用IDA Pro反汇编一个二进制文件的方法如下：\n1. 打开IDA Pro并选择“新建”。\n2. 选择适当的文件格式加载你的二进制文件。\n3. IDA Pro会自动分析二进制文件并提供反汇编视图。\n4. 你可以浏览反汇编的代码，以了解二进制文件的功能。\n5. 使用IDA Pro的交互功能重命名函数、添加注释，以便更容易分析。"
  ],
  "negative": [
    "使用IDA Pro进行文件反汇编的方法：\n1. 打开IDA Pro并选择“新建项目”。\n2. 加载任何类型的文件，IDA Pro会自动将其转换为源代码。\n3. 你可以直接运行反汇编代码，并通过调试器查看执行结果。\n4. 如果文件有加密，可以在IDA Pro中直接解密。\n5. 最后，生成一个全新的二进制文件。",
    "使用IDA Pro进行简单的文件修改：\n1. 打开IDA Pro并载入文件。\n2. 选择修改的部分并进行编辑。\n3. 保存修改后的文件。\n4. 测试修改后的文件是否工作正常。\n5. 完成所有修改后，生成新的文件。"
  ]
}

query是真实的用户咨询，LLM会提供两个答案，用户点赞选择的答案标记为positive，没有被选中的标记为negative！

1、先看embedding。训练样本的格式是[query、pos、neg]，微调的终极目的是让LLM的回答和query匹配，基于这个思路，设计出了Contrastive Learning，也叫Triplet Loss：先把三段文本求embedding，然后让query+pos的相似度最大，query+neg的相似度最小，loss的设计如下：

  ![外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传](https://img-home.csdnimg.cn/images/20230724024159.png?origin_url=https%3A%2F%2Fp0-xtjj-private.juejin.cn%2Ftos-cn-i-73owjymdk6%2F908a590d46ed41d1a43cdda709b623ba~tplv-73owjymdk6-jj-mark%3A0%3A0%3A0%3A0%3Aq75.awebp%3Fpolicy%3DeyJ2bSI6MywidWlkIjoiMzgxNzk2NzY5NjIyMTUzNCJ9%26rk3s%3De9ecf3d6%26x-orig-authkey%3Df32326d3454f2ac7e96d3d06cdbb035152127018%26x-orig-expires%3D1723024039%26x-orig-sign%3Dh1Y%252FuEj5sVa9yUoakS5StkS%252B3UA%253D&pos_id=img-q2Bi5rsJ-1722937705897)

q、p、n分别是三段text的embedding，d 是距离度量（例如欧氏距离或余弦相似度），α 是一个超参数，称为边际（margin）。这个loss函数意义直观，容易理解！具体怎么落地实现了？既然要计算相似度，那就干脆先把query和pos、neg的相似度事先全部先算好，放在矩阵里，便于后续取用。矩阵的每列都是用户每次反馈的数据。矩阵的第一列是query和pos的相似度，其他列是query和neg的相似度，如下：

sim_matrix = [[sim(q1, p1), sim(q1, n11), sim(q, n12), ...]
　　　　　　　　 [sim(q2, p2), sim(q2, n21), sim(q, n22), ...]]

因为第一列是query和pos的相似度，那么第一列的数值应该尽量大，其他列的数值应该尽量小，这不正好可以使用crossEntropy么？labels向量 = [1,0,0,0…]，经过crossEntropy相乘后，loss只剩query—pos的相似度啦！具体落地实现的方式稍微有些变通：

（1）以M3E微调为例，微调实现的代码在这里：https://github.com/wangyuxinwhy/uniem/blob/main/uniem/criteria.py#L62，核心的loss方法如下：

class TripletInBatchNegSoftmaxContrastLoss(ContrastLoss):
    def __init__(self, temperature: float = 0.05, add_swap_loss: bool = False):
        super().__init__(temperature)
        self.add_swap_loss = add_swap_loss
        if self.add_swap_loss:
            self._pair_contrast_softmax_loss = PairInBatchNegSoftmaxContrastLoss(temperature)
        else:
            self._pair_contrast_softmax_loss = None

    def forward(
        self,
        text_embeddings: torch.Tensor,
        text_pos_embeddings: torch.Tensor,
        text_neg_embeddings: torch.Tensor,
    ) -> torch.Tensor:
        # 计算正样本相似度向量
        sim_pos_vector = torch.cosine_similarity(text_embeddings, text_pos_embeddings, dim=-1)
        # 计算负样本相似度矩阵
        sim_neg_matrix = torch.cosine_similarity(
            text_embeddings.unsqueeze(1),
            text_neg_embeddings.unsqueeze(0),
            dim=-1,
        )
        # 将正样本相似度和负样本相似度拼接成一个矩阵
        sim_matrix = torch.cat([sim_pos_vector.unsqueeze(1), sim_neg_matrix], dim=1)
        # 温度缩放
        sim_matrix = sim_matrix / self.temperature
        # 生成标签，目的是让loss选择第一列的数值
        labels = torch.zeros(sim_matrix.size(0), dtype=torch.long, device=sim_matrix.device)
        # 计算交叉熵损失
        loss = torch.nn.CrossEntropyLoss()(sim_matrix, labels)
        # 如果有附加交换损失，则加上
        if self._pair_contrast_softmax_loss:
            loss += self._pair_contrast_softmax_loss(text_pos_embeddings, text_embeddings)
        return loss

uniem封装后，使用也很简单，几行代码就搞定了：

from datasets import load_dataset

from uniem.finetuner import FineTuner

dataset = load_dataset('/data/security_zh', 'STS-B')
# 指定训练的模型为 m3e-small
finetuner = FineTuner.from_pretrained('moka-ai/m3e-large', dataset=dataset)
finetuner.run(epochs=1)

M3E微调后的效果好不好，测评的方式有多种：

模型本身的指标：https://github.com/wangyuxinwhy/uniem/tree/main/mteb-zh 用文本分类、聚类、retrieve、rerank等方式
RAG的指标：https://www.cnblogs.com/theseventhson/p/18261594 context recall、context Precision
用户实际使用评价，核心还是triplet的点赞数据是不是够多

（2）同理，beg的baai_general_embedding微调的方法详见：https://github.com/FlagOpen/FlagEmbedding/blob/master/examples/finetune/README.md ；数据集格式如下，都是一样的：

{"query": str, "pos": List[str], "neg":List[str]}

重写getitem函数，

 def __getitem__(self, item) -> Tuple[str, List[str]]:
        query = self.dataset[item]['query']
        if self.args.query_instruction_for_retrieval is not None:
            query = self.args.query_instruction_for_retrieval + query

        passages = []

        assert isinstance(self.dataset[item]['pos'], list)
        pos = random.choice(self.dataset[item]['pos'])
        passages.append(pos)

        if len(self.dataset[item]['neg']) < self.args.train_group_size - 1:
            num = math.ceil((self.args.train_group_size - 1) / len(self.dataset[item]['neg']))
            negs = random.sample(self.dataset[item]['neg'] * num, self.args.train_group_size - 1)
        else:
            negs = random.sample(self.dataset[item]['neg'], self.args.train_group_size - 1)
        passages.extend(negs)

        if self.args.passage_instruction_for_retrieval is not None:
            passages = [self.args.passage_instruction_for_retrieval+p for p in passages]
        return query, passages

把原本的数据换个格式：

(
    "query：如何使用IDA Pro反汇编一个二进制文件？",
    [
        "passage：使用IDA Pro反汇编一个二进制文件的方法如下：\n1. 打开IDA Pro并选择“新建”。\n2. 选择适当的文件格式加载你的二进制文件。\n3. IDA Pro会自动分析二进制文件并提供反汇编视图。\n4. 你可以浏览反汇编的代码，以了解二进制文件的功能。\n5. 使用IDA Pro的交互功能重命名函数、添加注释，以便更容易分析。",
        "passage：使用IDA Pro进行文件反汇编的方法：\n1. 打开IDA Pro并选择“新建项目”。\n2. 加载任何类型的文件，IDA Pro会自动将其转换为源代码。\n3. 你可以直接运行反汇编代码，并通过调试器查看执行结果。\n4. 如果文件有加密，可以在IDA Pro中直接解密。\n5. 最后，生成一个全新的二进制文件。",
        "passage：IDA Pro是一款功能强大的反汇编工具，用户可以通过它轻松分析二进制文件。"
    ]
)

微调核心过程：

def encode(self, features):
        if features is None:
            return None
        psg_out = self.model(**features, return_dict=True)
        p_reps = self.sentence_embedding(psg_out.last_hidden_state, features['attention_mask'])
        if self.normlized:
            p_reps = torch.nn.functional.normalize(p_reps, dim=-1)
        return p_reps.contiguous()

    def compute_similarity(self, q_reps, p_reps):
        if len(p_reps.size()) == 2:
            return torch.matmul(q_reps, p_reps.transpose(0, 1))
        return torch.matmul(q_reps, p_reps.transpose(-2, -1))#矩阵相乘，本质还是内积

    def forward(self, query: Dict[str, Tensor] = None, passage: Dict[str, Tensor] = None, teacher_score: Tensor = None):
        q_reps = self.encode(query)
        p_reps = self.encode(passage)

        if self.training:
            if self.negatives_cross_device and self.use_inbatch_neg:
                q_reps = self._dist_gather_tensor(q_reps)
                p_reps = self._dist_gather_tensor(p_reps)

            group_size = p_reps.size(0) // q_reps.size(0)
            if self.use_inbatch_neg:
                scores = self.compute_similarity(q_reps, p_reps) / self.temperature # B B*G
                scores = scores.view(q_reps.size(0), -1)

                target = torch.arange(scores.size(0), device=scores.device, dtype=torch.long)
                target = target * group_size
                loss = self.compute_loss(scores, target)
            else:
                scores = self.compute_similarity(q_reps[:, None, :,], p_reps.view(q_reps.size(0), group_size, -1)).squeeze(1) / self.temperature # B G

                scores = scores.view(q_reps.size(0), -1)
                target = torch.zeros(scores.size(0), device=scores.device, dtype=torch.long)
                loss = self.compute_loss(scores, target)

        else:
            scores = self.compute_similarity(q_reps, p_reps)
            loss = None
        return EncoderOutput(
            loss=loss,
            scores=scores,
            q_reps=q_reps,
            p_reps=p_reps,
        )

    def compute_loss(self, scores, target):
        return self.cross_entropy(scores, target)

    def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
        if t is None:
            return None
        t = t.contiguous()

        all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
        dist.all_gather(all_tensors, t)

        all_tensors[self.process_rank] = t
        all_tensors = torch.cat(all_tensors, dim=0)

        return all_tensors

模型用的还是双塔结构 BiEncoderModel，先用矩阵相乘的形式得到query和passage中每条text的相似度，然后构造target向量，通过crossEntropy选择passage中的pos回答，这个落地实现的核心思路和M3E完全一样啊！微调的接口已经封装好了，直接调用：

torchrun \
> -m FlagEmbedding.baai_general_embedding.finetune.run \
> --output_dir /root/huggingface/bge_finetune \
> --model_name_or_path /root/huggingface/bge-large-zh-v1.5 \
> --train_data /root/huggingface/data/user_feedback \
> --learning_rate 1e-5 \
> --num_train_epochs 5 \
> --dataloader_drop_last True \
> --normlized True \
> --temperature 0.02 \
> --query_max_len 64 \
> --passage_max_len 256 \
> --train_group_size 2 \
> --negatives_cross_device \
> --logging_steps 10 \
> --save_steps 1000

运行完毕：

在这里插入图片描述

微调数据量有限的情况下，epoch越多，loss越小！

在这里插入图片描述

epoche=30，loss降至0.999

在这里插入图片描述

epoche=60，loss降至0.499；

在这里插入图片描述

微调完后测评的脚本也是现成的：https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/baai_general_embedding/finetune/eval_msmarco.py 核心思路是对query做encode，然后查找100个最接近的answer，然后计算Recall和MRR；可以直接执行命令：

python -m FlagEmbedding.baai_general_embedding.finetune.eval_msmarco \
--encoder /root/huggingface/bge-large-zh-v1.5 \
--fp16 \
--add_instruction \
--k 100 \
--corpus_data /root/huggingface/data/sec_corpus.json \
--query_data /root/huggingface/data/sec_query.json

corpus_data包含了想要检索的内容：

{"content": "为什么要对数据加读锁了而不是互斥锁了？在互斥机制中，读者和写者都需要独立独占互斥量以独占共享资源；而在读写锁机制下，允许同时有多个读者读访问共享资源，只有写者才需要独占资源。相比互斥机制，读写机制由于允许多个读者同时读访问共享资源，进一步提高了多线程的并发度"}
{"content": "从R4+C的地方取4字节数据存入R0，然后把R0存到栈上；接着把R0+4，这里ida已经识别出了是rwlock读写锁，然后就是调用pthread_rwlock_rdlock获取读写锁的读锁！这就很关键了"}
{"content": "从ida的trace记录看，前面所有的指令都没有读取栈上保存的url+http头的数据，所以前面肯定还没来得及生成那4个加密字段；从这里开始用读写锁，结合上面的分析大胆猜测：接下来要开始生成加密字段了！"}
{"content": "第三个参数我是用frida hook得到的，换了个环境地址肯定也变了，所以这里直接”抄袭“拿过来用肯定报错，这种反调试的方法实在是秒啊！动态调试暂时卡壳"}
{"content": " 之前通过hook registerNative发现：metasec_ml中的0x1094c被ms.bd.c.h.a方法注册成了native函数，这是metasec_ml唯一的native函数，肯定很重要，就从这里下手呗！这个函数有5个参数，分别都是啥了？"}

query_data包含了问题和正确答案，如下：

{"query": "frida是什么？", "positive": ["Frida是一款基于python + javascript 的hook框架，适用于android/ios/linux/win/osx等平台。Frida的动态代码执行功能，主要是在它的核心引擎Gum中用C语言来实现的", "只要兼容V8引擎就能正常使用frida"]}
{"query": "怎么使用IDA?", "positive": ["1、安装IDA   2、用IDA打开二进制文件，可以使用F5将汇编反编译成C语言伪代码   3、可以直接调试伪代码了解二进制代码逻辑"]}
{"query": "怎么脱壳?", "positive": ["对于一代、二代壳，可以直接使用frida dexdump从内存把正常的dex代码dump到磁盘"]}

从测评的原理来看，和https://www.cnblogs.com/theseventhson/p/18261594 这里面对整个RAG评测是一样的，所以直接采取RAG的评测方法！

2、reranker微调，这里以beg的reranker为例：https://github.com/FlagOpen/FlagEmbedding/blob/master/examples/reranker/README.md ；训练样本的格式和embedding是一样的，但是也要先对训练样本的格式做转换：

def __getitem__(self, item) -> List[BatchEncoding]:
    # 获取当前数据项的 query 和正样本
    query = self.dataset[item]['query']
    pos = random.choice(self.dataset[item]['pos'])

    # 如果负样本数量不足，则重复采样
    if len(self.dataset[item]['neg']) < self.args.train_group_size - 1:
        num = math.ceil((self.args.train_group_size - 1) / len(self.dataset[item]['neg']))
        negs = random.sample(self.dataset[item]['neg'] * num, self.args.train_group_size - 1)
    else:
        # 随机选择 train_group_size - 1 个负样本
        negs = random.sample(self.dataset[item]['neg'], self.args.train_group_size - 1)

    # 初始化批次数据列表
    batch_data = []
    
    # 添加正样本
    batch_data.append(self.create_one_example(query, pos))
    
    # 添加负样本
    for neg in negs:
        batch_data.append(self.create_one_example(query, neg))

    return batch_data  # 返回正负样本组合的批次数据

batch_data前面是pos样本，后面接着neg样本，每个batch_data的格式如下：

batch_data = [
    BatchEncoding({
        'input_ids': [101, ...],       # pos 编码后的 token ID
        'attention_mask': [1, 1, ...]  # 注意力掩码
    }),
    BatchEncoding({
        'input_ids': [101, ...],       # neg 编码后的 token ID
        'attention_mask': [1, 1, ...]  # 注意力掩码
    }),
    BatchEncoding({
        'input_ids': [101, ...],       # neg 编码后的 token ID
        'attention_mask': [1, 1, ...]  # 注意力掩码
    })
    .......
]

底层本质还是个分类模型，使用的是SequenceClassifierOutput

class CrossEncoder(nn.Module):
    def __init__(self, hf_model: PreTrainedModel, model_args: ModelArguments, data_args: DataArguments,
                 train_args: TrainingArguments):
        super().__init__()
        self.hf_model = hf_model
        self.model_args = model_args
        self.train_args = train_args
        self.data_args = data_args

        self.config = self.hf_model.config
        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')

        self.register_buffer(
            'target_label',
            torch.zeros(self.train_args.per_device_train_batch_size, dtype=torch.long)
        )

    def gradient_checkpointing_enable(self, **kwargs):
        self.hf_model.gradient_checkpointing_enable(**kwargs)

    def forward(self, batch):
        #选择分类模型
        ranker_out: SequenceClassifierOutput = self.hf_model(**batch, return_dict=True)
        logits = ranker_out.logits

        if self.training:
            scores = logits.view(
                self.train_args.per_device_train_batch_size,
                self.data_args.train_group_size
            )
            #通过target_label选择pos列用于计算loss的分母
            loss = self.cross_entropy(scores, self.target_label)

            return SequenceClassifierOutput(
                loss=loss,#输入loss反向传播更新参数
                **ranker_out,
            )
        else:
            return ranker_out

    @classmethod
    def from_pretrained(
            cls, model_args: ModelArguments, data_args: DataArguments, train_args: TrainingArguments,
            *args, **kwargs
    ):
        hf_model = AutoModelForSequenceClassification.from_pretrained(*args, **kwargs)
        reranker = cls(hf_model, model_args, data_args, train_args)
        return reranker

    def save_pretrained(self, output_dir: str):
        state_dict = self.hf_model.state_dict()
        state_dict = type(state_dict)(
            {k: v.clone().cpu()
             for k,
             v in state_dict.items()})
        self.hf_model.save_pretrained(output_dir, state_dict=state_dict)

如何系统的去学习AI大模型LLM ？

作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。

但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的 AI大模型资料 包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

所有资料 ⚡️ ，朋友们如果有需要全套《LLM大模型入门+进阶学习资源包》，扫码获取~

👉CSDN大礼包🎁：全网最全《LLM大模型入门+进阶学习资源包》免费分享（安全链接，放心点击）👈

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。那以下这些PDF籍就是非常不错的学习资源。

在这里插入图片描述

四、AI大模型商业化落地方案

阶段1：AI大模型时代的基础理解

目标：了解AI大模型的基本概念、发展历程和核心原理。
内容：
- L1.1 人工智能简述与大模型起源
- L1.2 大模型与通用人工智能
- L1.3 GPT模型的发展历程
- L1.4 模型工程
  - L1.4.1 知识大模型
  - L1.4.2 生产大模型
  - L1.4.3 模型工程方法论
  - L1.4.4 模型工程实践
- L1.5 GPT应用案例

阶段2：AI大模型API应用开发工程

目标：掌握AI大模型API的使用和开发，以及相关的编程技能。
内容：
- L2.1 API接口
  - L2.1.1 OpenAI API接口
  - L2.1.2 Python接口接入
  - L2.1.3 BOT工具类框架
  - L2.1.4 代码示例
- L2.2 Prompt框架
  - L2.2.1 什么是Prompt
  - L2.2.2 Prompt框架应用现状
  - L2.2.3 基于GPTAS的Prompt框架
  - L2.2.4 Prompt框架与Thought
  - L2.2.5 Prompt框架与提示词
- L2.3 流水线工程
  - L2.3.1 流水线工程的概念
  - L2.3.2 流水线工程的优点
  - L2.3.3 流水线工程的应用
- L2.4 总结与展望

阶段3：AI大模型应用架构实践

目标：深入理解AI大模型的应用架构，并能够进行私有化部署。
内容：
- L3.1 Agent模型框架
  - L3.1.1 Agent模型框架的设计理念
  - L3.1.2 Agent模型框架的核心组件
  - L3.1.3 Agent模型框架的实现细节
- L3.2 MetaGPT
  - L3.2.1 MetaGPT的基本概念
  - L3.2.2 MetaGPT的工作原理
  - L3.2.3 MetaGPT的应用场景
- L3.3 ChatGLM
  - L3.3.1 ChatGLM的特点
  - L3.3.2 ChatGLM的开发环境
  - L3.3.3 ChatGLM的使用示例
- L3.4 LLAMA
  - L3.4.1 LLAMA的特点
  - L3.4.2 LLAMA的开发环境
  - L3.4.3 LLAMA的使用示例
- L3.5 其他大模型介绍

阶段4：AI大模型私有化部署

目标：掌握多种AI大模型的私有化部署，包括多模态和特定领域模型。
内容：
- L4.1 模型私有化部署概述
- L4.2 模型私有化部署的关键技术
- L4.3 模型私有化部署的实施步骤
- L4.4 模型私有化部署的应用场景

学习计划：

阶段1：1-2个月，建立AI大模型的基础知识体系。
阶段2：2-3个月，专注于API应用开发能力的提升。
阶段3：3-4个月，深入实践AI大模型的应用架构和私有化部署。
阶段4：4-5个月，专注于高级模型的应用和部署。

这份完整版的所有 ⚡️ 大模型 LLM 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【`保证100%免费`】

全套《LLM大模型入门+进阶学习资源包》↓↓↓ 获取~

👉CSDN大礼包🎁：全网最全《LLM大模型入门+进阶学习资源包》免费分享（安全链接，放心点击）👈

程序学到昏

关注

9
点赞
踩
10

收藏

觉得还不错? 一键收藏
0
评论
LLM大模型: RAG两大核心利器 — embedding和reranker模型微调fine-tune

要想RAG好，embedding和reranker必须给力！目前市面上流行的embedding和reranker使用的都是通用预料训练，并未针对安全这个细分领域定制，所以有必要使用安全领域的预料微调一下！
复制链接

扫一扫