bert_cate/other/main.py

27 lines
745 B
Python

from transformers import pipeline
classifier = pipeline(
"zero-shot-classification",
model="./nlp_structbert_zero-shot-classification_chinese-base",
device="cpu",
max_length=512,
ignore_mismatched_sizes=True # 忽略维度不匹配警告
)
level1 = [
"食品", "电器", "洗护", "女装", "手机",
"健康", "男装", "美妆", "电脑", "运动",
"内衣", "母婴", "数码", "百货", "鞋包",
"办公", "家装", "饰品", "车品", "图书",
"生鲜", "家纺", "宠物", "奢品", "其它"
]
result = classifier(
"CONBA/康恩贝森兰康牌铁皮石斛西洋参颗粒 规格:3g/包*60包:",
candidate_labels=level1,
truncation=True # 显式指定截断策略
)
print(result)