如何基于多分类数据集生成三元组并微调SentenceTransformer？

核心逻辑 #

• 正三元组：从同一类别中选取两个不同文本，标记为1（代表语义相似）
• 负三元组：从不同类别中各选一个文本，标记为0（代表语义不相似）
• 平衡策略：控制正负样本比例（常用1:1或1:2），避免某类样本过多导致模型偏倚

代码实现 #

假设你的分类数据集是包含text（文本）和label（类别）列的DataFrame，以下是直接可用的生成逻辑：

import pandas as pd
import random
from tqdm import tqdm

def generate_triples(df, pos_neg_ratio=1):
    # 按类别分组，方便后续采样
    grouped = df.groupby('label')
    class_texts = {label: group['text'].tolist() for label, group in grouped}
    all_classes = list(class_texts.keys())
    triples = []

    # 生成正三元组：同一类别内的文本对
    for label, texts in class_texts.items():
        text_num = len(texts)
        if text_num < 2:
            continue  # 类别下文本数不足2，跳过生成正样本
        # 按比例生成对应数量的正样本
        for _ in range(text_num * pos_neg_ratio):
            idx1, idx2 = random.sample(range(text_num), 2)
            triples.append((texts[idx1], texts[idx2], 1))

    # 生成负三元组：数量匹配正样本
    pos_sample_count = len(triples)
    for _ in tqdm(range(pos_sample_count)):
        # 随机选两个不同类别
        label_a, label_b = random.sample(all_classes, 2)
        text_a = random.choice(class_texts[label_a])
        text_b = random.choice(class_texts[label_b])
        triples.append((text_a, text_b, 0))

    # 打乱样本顺序
    random.shuffle(triples)
    return triples

# 示例用法
if __name__ == "__main__":
    # 模拟多分类数据集
    sample_df = pd.DataFrame({
        'text': [
            "猫喜欢吃鱼", "猫咪怕水", "狗狗爱啃骨头", "小狗喜欢跑跳",
            "苹果是水果", "香蕉富含钾", "汽车有四个轮子", "卡车用来拉货"
        ],
        'label': ['猫', '猫', '狗', '狗', '水果', '水果', '交通工具', '交通工具']
    })
    # 生成三元组
    train_triples = generate_triples(sample_df)
    # 打印前5个样本
    for item in train_triples[:5]:
        print(item)

优化技巧 #

• 硬负样本挖掘：如果模型效果一般，可替换随机负样本为「硬负样本」——用预训练SentenceTransformer计算文本相似度，挑选与正样本语义接近但属于不同类别的文本，能提升模型区分度
• 动态采样：若数据集过大，可在训练循环中动态生成三元组，避免一次性生成大量数据占用内存
• 类别均衡处理：针对类别不平衡的数据集，对小类别多采样正样本，对大类别限制采样数量，保证每个类别的样本贡献均衡

内容的提问来源于stack exchange，提问作者Yogesh Haribhau Kulkarni