要想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完全一样啊！微调完后测评的脚本也是现成的：https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/baai_general_embedding/finetune/eval_msmarco.py 核心思路是对query做encode，然后查找100个最接近的answer，然后计算Recall和MRR
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介绍与使用入门
标签：RAG,loss,模型,args,reps,embedding,query,self,size From： https://www.cnblogs.com/theseventhson/p/18288053