Transformers Pipeline: 仅需几行代码运行 NLP 模型
1. 总览
Hugging Face Transformers 提供了 pipeline() API,让你只需几行代码就能运行经过预训练的 NLP/vision/audio 模型,无需手动加载 tokenizer、config 或处理 tensor。本文帮助你理解 pipeline() 的工作原理,如何在常见任务中使用它(情感分析、NER、摘要、零样本分类),以及何时需要“深入底层”来获得更多的控制权。
读完本文后,你可以在自己的 Python 应用中运行 Hugging Face 模型推理、编写测试来检查输出,以及在投入生产前了解 pipeline 的限制。
2. 安装
你需要 transformers 库和一个后端(PyTorch 或 TensorFlow)。例如使用 PyTorch:
pip install transformers torch
推荐版本(撰写本文时):transformers>=4.40,torch>=2.2。
注意: 首次运行时,
pipeline()会从 Hugging Face Hub 下载模型并缓存到~/.cache/huggingface。模型可能有数百 MB,所以如果在容器 CI 中运行,你需要为这一步计算时间/网络开销,或者预先将缓存烘焙到镜像中。
3. Pipeline 基础:情感分析
from transformers import pipeline
classifier = pipeline("sentiment-analysis")
result = classifier("I really love this product, it works great!")
print(result)
- 第 1 行:导入
pipeline函数。 - 第 3 行:为
sentiment-analysis任务创建 pipeline。当不指定模型时,Hugging Face 会自动选择默认模型(通常是distilbert-base-uncased-finetuned-sst-2-english)。 - 第 4 行:像调用函数一样调用 pipeline,将文本句子传入。
输出:
[{'label': 'POSITIVE', 'score': 0.9998}]
pipeline 返回一个列表,因为它支持批量处理(传入字符串列表也可以)。
4. 指定具体的模型和 Tokenizer
不应该在生产代码中依赖默认模型,因为 Hugging Face 可能在版本间改变默认模型。应该明确指定:
from transformers import pipeline
classifier = pipeline(
task="sentiment-analysis",
model="distilbert-base-uncased-finetuned-sst-2-english",
device=-1 # -1 = CPU,0 = 第一个 GPU
)
results = classifier([
"This movie was a complete waste of time.",
"Absolutely fantastic service!"
])
for r in results:
print(r)
model:明确指定 checkpoint,避免依赖可能变化的默认值。device:-1在 CPU 上运行,0在 GPU 索引 0 上运行。如果机器没有 CUDA 但你设置device=0,pipeline 会抛出错误。- 将文本列表传入 classifier 来处理批量,这比逐个调用更高效。
输出:
{'label': 'NEGATIVE', 'score': 0.9995}
{'label': 'POSITIVE', 'score': 0.9997}
5. 命名实体识别(NER)
from transformers import pipeline
ner = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
text = "Satya Nadella is the CEO of Microsoft, based in Redmond."
for entity in ner(text):
print(entity["word"], "->", entity["entity_group"], f"{entity['score']:.2f}")
aggregation_strategy="simple":合并子 token(例如Sat、##ya)成为完整的实体,而不是返回零散的 token 片段。entity_group:实体类型标签(PER、ORG、LOC 等)。
输出:
Satya Nadella -> PER 0.99
Microsoft -> ORG 0.99
Redmond -> LOC 0.98
注意: 如果省略
aggregation_strategy,你会收到单个离散 token 和 BIO 标签(B-PER、I-PER等),需要手动合并——如果处理不当容易引入 bug。
6. 零样本分类
零样本允许将文本分类到任意标签,模型从未专门针对这些标签进行训练。
from transformers import pipeline
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
result = classifier(
"This new phone has an amazing camera and battery life.",
candidate_labels=["technology", "sports", "politics", "finance"]
)
print(result["labels"][0], result["scores"][0])
candidate_labels:你自定义的标签列表,无需对模型进行微调。- 模型使用自然语言推理技术来计算文本对每个标签的蕴含概率。
输出:
technology 0.93
零样本非常方便快速原型开发,但比直接微调的模型更慢、准确度更低。
7. 文本摘要
from transformers import pipeline
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
article = """
The Transformers library by Hugging Face provides thousands of pretrained
models to perform tasks on texts such as classification, information
extraction, question answering, summarization, translation, and more.
It bridges research and production by exposing a simple pipeline API
on top of complex model architectures.
"""
summary = summarizer(article, max_length=40, min_length=10, do_sample=False)
print(summary[0]["summary_text"])
max_length/min_length:限制摘要的 token 长度,不是单词数。do_sample=False:使用贪心/束搜索解码而不是随机采样,产生更稳定的结果。
输出(缩写):
Transformers library provides pretrained models for classification,
extraction, question answering, and summarization.
8. Pipeline 内部原理
pipeline() 实际上拼接了 3 个组件:tokenizer、model 和针对每个任务的后处理逻辑。你可以自己组装来获得更多控制权:
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
model_name = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
inputs = tokenizer("Great product!", return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.softmax(logits, dim=-1)
label_id = probs.argmax().item()
print(model.config.id2label[label_id], probs[0][label_id].item())
AutoTokenizer/AutoModelForSequenceClassification:根据模型名称自动识别正确的架构。tokenizer(..., return_tensors="pt"):返回 PyTorch tensor 而不是普通列表。torch.no_grad():禁用梯度计算,因为这是推理而不是训练——有助于节省内存并提高速度。model.config.id2label:将类索引映射到标签名称(LABEL_0/LABEL_1或POSITIVE/NEGATIVE等)。
理解这一层有助于你调试 pipeline 返回异常结果的情况,或当需要自定义逻辑时(例如应用自定义阈值而不是取 argmax)。
9. 为 Pipeline 编写测试
由于模型输出在版本间可能略有变化,测试应该检查“行为”而不是比较绝对的小数值。
import pytest
from transformers import pipeline
@pytest.fixture(scope="module")
def sentiment_classifier():
return pipeline(
"sentiment-analysis",
model="distilbert-base-uncased-finetuned-sst-2-english",
)
def test_classifies_positive_sentence(sentiment_classifier):
result = sentiment_classifier("I absolutely love this!")[0]
assert result["label"] == "POSITIVE"
assert result["score"] > 0.9
def test_classifies_negative_sentence(sentiment_classifier):
result = sentiment_classifier("This is the worst experience ever.")[0]
assert result["label"] == "NEGATIVE"
assert result["score"] > 0.9
scope="module":fixture 只为整个测试文件加载一次模型,避免在每个测试用例中重复加载——加载模型需要数秒到数十秒。- 基于
label和score阈值进行断言,不硬编码精确的 score 值,因为它在不同库版本间可能略有偏差。
注意: 运行此类测试需要加载真实模型,因此不适合作为单元测试在快速 CI 循环中运行。应该将其分离为集成测试,或在测试周围业务逻辑时 mock pipeline。
10. 何时不应该使用 Pipeline
Pipeline 对于原型开发和标准任务非常好用,但有限制:
- 大批量、高性能场景: pipeline 的优化不如自己管理批处理、padding 和动态批处理。
- 需要复杂的自定义逻辑: 例如应用多个后处理条件、结合多个模型或访问隐藏状态——此时应该使用第 8 节介绍的
AutoModel+ tokenizer。 - 大规模生产部署: 应该考虑专门的工具,如 Text Generation Inference 或 ONNX Runtime/TensorRT 来优化延迟,而不是在请求处理程序中直接调用
pipeline()。 - 对 CPU 来说太大的模型: bart-large-mnli 或具有数十亿参数的 LLM 在 CPU 上运行会非常慢;需要 GPU 或量化。
11. 总结
Hugging Face Transformers 的 pipeline() 是运行预训练 NLP 模型推理的最快方法——用于情感分析、NER、摘要或零样本分类只需几行代码。当需要对性能或处理逻辑进行更深入的控制时,你可以拆解 pipeline 成单独的 tokenizer 和模型来自己管理。始终明确指定模型而不是使用默认值,编写基于行为(标签、score 阈值)的测试而不是绝对值,以避免库更新时测试不稳定。