如何用SetFit模型基于双文本输入实现少样本文本分类

SetFit默认仅支持单文本输入，但可以通过文本拼接或自定义嵌入融合的方式处理A、B两列文本输入，实现对C列标签的预测，以下是两种可行方案：

方案一：文本拼接（简单高效，推荐） #

直接将A、B两列的文本内容拼接为一个新文本列，作为SetFit的输入特征，无需修改模型结构，是最快捷的实现方式。

步骤示例 #

1. 预处理DataFrame，拼接文本列

import pandas as pd

# 假设你的数据集存储在df中
df["combined_text"] = df["A"] + " [SEP] " + df["B"]  # 用[SEP]分隔两列内容，帮助模型区分不同输入源

2. 加载SetFit模型并训练

from setfit import SetFitModel, SetFitTrainer
from sentence_transformers.losses import CosineSimilarityLoss

# 加载预训练的Sentence-BERT骨干模型
model = SetFitModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# 初始化训练器
trainer = SetFitTrainer(
    model=model,
    train_dataset=df[["combined_text", "C"]],
    loss_class=CosineSimilarityLoss,
    batch_size=16,
    num_epochs=1,
    num_iterations=20,  # SetFit对比学习的迭代次数
    column_mapping={"combined_text": "text", "C": "label"}  # 映射列名到SetFit预期的字段名
)

# 启动训练
trainer.train()

3. 预测新数据

# 对新数据先执行同样的拼接操作
new_data = pd.DataFrame({"A": ["新输入文本A"], "B": ["新输入文本B"]})
new_data["combined_text"] = new_data["A"] + " [SEP] " + new_data["B"]

# 生成预测结果
predictions = trainer.model.predict(new_data["combined_text"])

方案二：自定义嵌入融合（进阶，精细控制） #

如果需要更灵活地融合两列文本的嵌入特征，可以修改SetFit的模型结构，将A、B列的嵌入向量拼接后再输入分类头。

步骤示例 #

1. 自定义双输入SetFit模型

from setfit import SetFitModel
import torch.nn as nn

# 加载预训练的Sentence-BERT骨干模型
backbone = SetFitModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2").body

# 定义双输入融合模型：拼接两列文本的嵌入后接分类头
class DualInputSetFitModel(nn.Module):
    def __init__(self, backbone, num_labels):
        super().__init__()
        self.backbone = backbone
        emb_dim = backbone.get_sentence_embedding_dimension()
        self.classifier = nn.Linear(emb_dim * 2, num_labels)
    
    def forward(self, text_a, text_b):
        emb_a = self.backbone(text_a)["sentence_embedding"]
        emb_b = self.backbone(text_b)["sentence_embedding"]
        combined_emb = torch.cat([emb_a, emb_b], dim=1)
        return self.classifier(combined_emb)

# 初始化自定义模型
num_labels = df["C"].nunique()
model = DualInputSetFitModel(backbone, num_labels)

2. 自定义训练逻辑
   这种方式需要手动编写PyTorch训练循环，或修改SetFitTrainer的训练流程，适合具备PyTorch基础、需要极致定制的用户。若无需特殊定制，方案一已足够满足需求。

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