HuggingFace | 基础组件之Pipeline
什么是Pipeline
- Pipeline
- 将数据预处理、模型调用、结果后处理三部分组装成的流水线
- 使我们能够直接输入文本便获得最终的答案
Pipeline支持的任务类型
查看Pipeline支持的任务类型:
from transformers.pipelines import SUPPORTED_TASKS
print(SUPPORTED_TASKS.items())
for k, v in SUPPORTED_TASKS.items():
print(k, v)
通过SUPPORTED_TASKS
可以查看可以支持的任务。可以查看各种任务的实现类是什么,tf实现的类和pt实现的类。
Pipeline的创建与使用方式
- 根据任务类型直接创建Pipeline
pipe = pipeline("text-classification")
- 指定任务类型,再指定模型,创建基于指定模型的Pipeline
pipe = pipeline("text-classification", model="uer/roberta-base-finetuned-dianping-chinese")
- 预先加载模型,再创建Pipeline
# 这种方式,必须同时指定model和tokenizer
model = AutoModelForSequenceClassification.from_pretrained("uer/roberta-base-finetuned-dianping-chinese")
tokenizer = AutoTokenizer.from_pretrained("uer/roberta-base-finetuned-dianping-chinese")
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
- 使用GPU进行推理加速
pipe = pipeline("text-classification", model="uer/roberta-base-finetuned-dianping-chinese", device=0)
测试CPU和GPU运行的速度:
pipe.model.device
# device(type='cuda', index=0)
import torch
import time
times = []
for i in range(100):
torch.cuda.synchronize()
start = time.time()
pipe("我觉得不太行!")
torch.cuda.synchronize()
end = time.time()
times.append(end - start)
print(sum(times) / 100)
# cpu时间:2.1s
# gpu时间:0.6s
Pipeline的背后实现
- Step1 初始化Tokenizer
tokenizer = AutoTokenizer.from_pretrained("uer/roberta-base-finetuned-dianping-chinese")
- Step2 初始化Model
model = AutoModelForSequenceClassification.from_pretrained("uer/roberta-base-finetuned-dianping-chinese")
- Step3 数据预处理
input_text = "我觉得不太行!"
inputs = tokenizer(input_text, return_tensors="pt")
- Step4 模型预测
res = model(**inputs)
logits = res.logits
- Step5 结果后处理
logits = torch.softmax(logits, dim=-1)
pred = torch.argmax(logits).item()
result = model.config.id2label.get(pred)
代码整合:
# 导入依赖
from transformers import *
import torch
# 初始化模型 和 初始化Tokenizer
tokenizer = AutoTokenizer.from_pretrained("uer/roberta-base-finetuned-dianping-chinese")
model = AutoModelForSequenceClassification.from_pretrained("uer/roberta-base-finetuned-dianping-chinese")
# 数据预处理
input_text = "我觉得不太行!"
inputs = tokenizer(input_text, return_tensors="pt")
# 模型预测
res = model(**inputs)
logits = res.logits
# 查看标签属性
model.config.id2label # {0: 'negative (stars 1, 2 and 3)', 1: 'positive (stars 4 and 5)'}
# 结果后处理
logits = torch.softmax(logits, dim=-1)
pred = torch.argmax(logits).item()
result = model.config.id2label.get(pred) # 'negative (stars 1, 2 and 3)'
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· DeepSeek “源神”启动!「GitHub 热点速览」
· 微软正式发布.NET 10 Preview 1:开启下一代开发框架新篇章
· 我与微信审核的“相爱相杀”看个人小程序副业
· C# 集成 DeepSeek 模型实现 AI 私有化(本地部署与 API 调用教程)
· spring官宣接入deepseek,真的太香了~