训练与微调多向量嵌入模型

pip install -U sentence-transformers torch datasets

from sentence_transformers import SentenceTransformer, models
import torch.nn as nn

# 1. 加载基础 Transformer 模型
word_embedding_model = models.Transformer("bert-base-uncased", max_seq_length=256)

# 2. 添加投影层以减少每个 Token 的嵌入大小（例如降至 128 维）
# 这对于控制生产环境中的索引大小至关重要。
class TokenProjection(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim, bias=False)
        
    def forward(self, features):
        # 形状: [batch_size, seq_length, input_dim]
        token_embeddings = features["token_embeddings"]
        projected = self.linear(token_embeddings)
        # 归一化嵌入以进行点积搜索
        normalized = projected / projected.norm(dim=-1, keepdim=True)
        features.update({"token_embeddings": normalized})
        return features

# 实例化投影层
projection_layer = models.Dense(
    in_features=word_embedding_model.get_word_embedding_dimension(),
    out_features=128,
    activation_function=None,
    bias=False
)

# 组合成完整的模型结构
# 注意：我们这里不添加 Pooling 层，从而保留了 Token 级别的输出。
model = SentenceTransformer(modules=[word_embedding_model, projection_layer])

from datasets import Dataset

# 示例数据集结构
training_data = [
    {
        "query": "多向量嵌入是如何工作的？",
        "positive": "多向量嵌入将文档表示为一系列 Token 向量，并在检索时通过后期交互计算相似度。",
        "negative": "单向量嵌入将整个文本压缩为一个向量，这可能会丢失具体的 Token 细节。"
    },
    {
        "query": "后期交互有什么好处？",
        "positive": "后期交互允许预先计算文档向量，同时在查询时保持 Token 级别的对齐评分。",
        "negative": "早期交互模型将查询和文档一起处理，速度较慢，且无法进行预索引。"
    }
]

dataset = Dataset.from_list(training_data)

import torch
import torch.nn as nn

class LateInteractionTripletLoss(nn.Module):
    def __init__(self, model, margin=0.2):
        super().__init__()
        self.model = model
        self.margin = margin

    def maxsim(self, query_embeddings, doc_embeddings, query_mask, doc_mask):
        # query_embeddings: [B, M, D]
        # doc_embeddings: [B, N, D]
        # 计算相似度矩阵: [B, M, N]
        sim_matrix = torch.matmul(query_embeddings, doc_embeddings.transpose(1, 2))
        
        # 掩盖 Padding 填充 Token
        mask = doc_mask.unsqueeze(1) * query_mask.unsqueeze(2)
        sim_matrix = sim_matrix.masked_fill(mask == 0, -1e9)
        
        # MaxSim 计算: 沿文档 Token 取最大值，沿查询 Token 求和
        max_sim_per_query_token, _ = sim_matrix.max(dim=2)
        # 清零被掩盖的查询 Token
        max_sim_per_query_token = max_sim_per_query_token * query_mask
        return max_sim_per_query_token.sum(dim=1)

    def forward(self, sentence_features, labels):
        # 提取查询、正例和反例的特征
        reps = [self.model(x) for x in sentence_features]
        q_rep, p_rep, n_rep = reps[0], reps[1], reps[2]
        
        # 提取 Token 嵌入和注意力掩码
        q_emb, p_emb, n_emb = q_rep["token_embeddings"], p_rep["token_embeddings"], n_rep["token_embeddings"]
        q_mask, p_mask, n_mask = q_rep["attention_mask"], p_rep["attention_mask"], n_rep["attention_mask"]
        
        # 计算得分
        pos_scores = self.maxsim(q_emb, p_emb, q_mask, p_mask)
        neg_scores = self.maxsim(q_emb, n_emb, q_mask, n_mask)
        
        # 三元组损失
        loss = torch.clamp(self.margin - pos_scores + neg_scores, min=0.0)
        return loss.mean()

from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments

# 配置训练参数
args = SentenceTransformerTrainingArguments(
    output_dir="multi-vector-model",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,
    logging_steps=10,
)

# 初始化自定义损失
loss_fn = LateInteractionTripletLoss(model=model)

# 创建训练器
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    loss=loss_fn,
)

# 开始训练
trainer.train()

# 保存微调后的模型
model.save_pretrained("fine-tuned-colbert-model")