如何基于Transformers加速5000条数据集的零样本分类？

你的代码效率低的核心原因是逐条调用模型，没有利用批量推理的并行能力，以下是基于Transformers和Datasets的优化方案：

1. 直接使用Pipeline的批量处理能力 #

Transformers的pipeline原生支持传入文本列表，模型会自动按最优批次处理，大幅减少单条调用的开销，GPU环境下提速效果更明显。

修改后的代码：

from transformers import pipeline

classifier = pipeline("zero-shot-classification", model='cross-encoder/nli-roberta-base', device=0)  # device=0 指定用GPU，无GPU可去掉
candidate_labels = ["Self-direction: action", "Achievement", "Security: personal", "Security: societal", "Benevolence: caring", "Universalism: concern"]

# 提取所有文本到列表
texts = reduced_dataset["text"].tolist()

# 批量推理，可手动指定batch_size（如batch_size=32）适配硬件
results = classifier(texts, candidate_labels)

# 用列表推导式简化0/1标签转换
predictions = [[1 if score >= 0.5 else 0 for score in res["scores"]] for res in results]

2. 用Datasets库的map函数批量处理 #

如果你的数据集可以转为Datasets格式，map函数的批量处理能力能进一步优化效率，还支持多进程加速（CPU环境适用）。

代码示例：

from transformers import pipeline
from datasets import Dataset

classifier = pipeline("zero-shot-classification", model='cross-encoder/nli-roberta-base', device=0)
candidate_labels = ["Self-direction: action", "Achievement", "Security: personal", "Security: societal", "Benevolence: caring", "Universalism: concern"]

# Pandas DataFrame转Dataset
dataset = Dataset.from_pandas(reduced_dataset)

def batch_process(batch):
    results = classifier(batch["text"], candidate_labels)
    batch["predictions"] = [[1 if s >= 0.5 else 0 for s in res["scores"]] for res in results]
    return batch

# batched=True开启批量处理，num_proc是多进程数（GPU环境建议设为1）
processed_dataset = dataset.map(batch_process, batched=True, batch_size=32, num_proc=2)

# 提取最终结果
predictions = processed_dataset["predictions"]

3. 手动加载模型与分词器，精细化控制批量流程 #

如果需要更灵活的调整（如自定义分词、适配大批次），可以手动加载模型组件，完全控制推理流程：

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model_name = 'cross-encoder/nli-roberta-base'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

candidate_labels = ["Self-direction: action", "Achievement", "Security: personal", "Security: societal", "Benevolence: caring", "Universalism: concern"]
texts = reduced_dataset["text"].tolist()

# cross-encoder需要生成所有文本-标签对
all_pairs = [(text, label) for text in texts for label in candidate_labels]

# 批量分词
inputs = tokenizer(
    [pair[0] for pair in all_pairs],
    [pair[1] for pair in all_pairs],
    truncation=True,
    padding=True,
    return_tensors="pt"
).to(device)

# 批量推理，关闭梯度计算加速
model.eval()
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits
    # 取蕴含类别的概率（cross-encoder NLI模型中logits[:,2]对应蕴含）
    scores = torch.softmax(logits, dim=1)[:, 2].cpu().numpy()

# 将分数按每个文本的6个标签分组，转换为0/1标签
scores_per_text = scores.reshape(len(texts), len(candidate_labels))
predictions = (scores_per_text >= 0.5).astype(int).tolist()

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