LLM大模型: RAG两大核心利器 — embedding和reranker模型微调fine-tune
要想RAG好，embedding和reranker必须给力！目前市面上流行的embedding和reranker使用的都是通用预料训练，并未针对安全这个细分领域定制，所以有必要使用安全领域的预料微调一下！目前所有的预料中，获取成本比较低、并且不需要专门投入人力标注的有两种：
- 网上各种安全论坛的博客、各大热门产品的漏洞说明等
- 用户的点赞反馈数据（chatGPT、copilot等都有该功能）
对于作者本人而言，用户的点赞反馈数据更容易获取，所以这里使用这类数据，借鉴RLHF-DPO的思路对embedding和reranker模型做微调！训练样本的数据格式如下：
{ "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的设计如下：
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)
参考：
1、https://www.bilibili.com/video/BV1bN4y1n7Ex/?spm_id_from=333.788&vd_source=241a5bcb1c13e6828e519dd1f78f35b2 https://github.com/yuanzhoulvpi2017/SentenceEmbedding 实现自己的sentence-embedding训练代码
2、https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune finetune the baai-general-embedding with your data.
3、https://github.com/FlagOpen/FlagEmbedding/blob/master/examples/reranker/README.md finetune the cross-encoder reranker with your data.
4、https://huggingface.co/moka-ai/m3e-base https://github.com/wangyuxinwhy/uniem https://github.com/wangyuxinwhy/uniem/blob/main/examples/finetune.ipynb M3E微调
5、https://www.cnblogs.com/xiaoqi/p/18034447/MTEB 搜索引擎RAG召回效果评测MTEB介绍与使用入门
