BGE-M3微调实战指南：从零开始构建高效语义检索模型

1次阅读

共计 3011 个字符，预计需要花费 8 分钟才能阅读完成。

语义检索模型在现代 NLP 应用中扮演着关键角色，尤其在需要处理复杂查询意图和多语言内容的场景中。与传统的基于关键词的检索系统相比，语义检索模型能够理解查询的深层含义，从而提供更精准的搜索结果。然而，构建高效的语义检索模型面临着多重挑战：

多语言支持：需要处理不同语言间的语义差异和资源不均衡问题
计算效率：大规模 embedding 生成和检索对计算资源要求高
数据质量：高质量的标注数据获取成本高昂
领域适应：预训练模型在特定领域的表现可能不足

BGE-M3 作为新一代语义检索模型，在传统 BERT 架构基础上进行了多项创新：

注意力机制改进：
采用稀疏注意力模式降低计算复杂度
引入跨层注意力共享机制
池化策略优化：
动态混合 CLS token 与平均池化
添加可学习的池化权重参数
多语言处理：
统一的子词切分策略
语言无关的位置编码

from torch.utils.data import Dataset
from typing import Dict, List
import json

class RetrievalDataset(Dataset):
    def __init__(self, file_path: str, tokenizer, max_length: int = 512):
        self.data = []
        with open(file_path) as f:
            for line in f:
                item = json.loads(line)
                self.data.append(item)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:
        item = self.data[idx]
        query = self.tokenizer(item["query"], 
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        passage = self.tokenizer(item["passage"],
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        return {"query_input_ids": query["input_ids"].squeeze(0),
            "query_attention_mask": query["attention_mask"].squeeze(0),
            "passage_input_ids": passage["input_ids"].squeeze(0),
            "passage_attention_mask": passage["attention_mask"].squeeze(0),
            "label": torch.tensor(item["label"], dtype=torch.float)
        }

import torch
import torch.nn as nn
import torch.nn.functional as F

class ContrastiveLossWithTemperature(nn.Module):
    def __init__(self, temp: float = 0.05):
        super().__init__()
        self.temp = temp
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, query_emb: torch.Tensor, passage_emb: torch.Tensor) -> torch.Tensor:
        # 计算相似度矩阵
        sim_matrix = torch.matmul(query_emb, passage_emb.T) / self.temp

        # 创建标签（对角线为正样本）labels = torch.arange(sim_matrix.size(0), device=query_emb.device)

        # 计算对比损失
        loss = (self.cross_entropy(sim_matrix, labels) + 
               self.cross_entropy(sim_matrix.T, labels)) / 2
        return loss

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()
accumulation_steps = 4

def train_step(model, batch, optimizer):
    with autocast():
        query_emb = model(input_ids=batch["query_input_ids"], 
                         attention_mask=batch["query_attention_mask"])
        passage_emb = model(input_ids=batch["passage_input_ids"], 
                           attention_mask=batch["passage_attention_mask"])
        loss = loss_fn(query_emb, passage_emb)

    # 梯度累积
    loss = loss / accumulation_steps
    scaler.scale(loss).backward()

    if (step + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

    return loss.item() * accumulation_steps