使用方法如下:
from transformers import BertTokenizer, BertForSequenceClassification
import torch
# 加载已训练的模型和分词器
model_path = 'qixun/tangsong_poem_classify'
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path)
# 预处理函数
def preprocess_text(text):
inputs = tokenizer(text, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
return inputs
# 分类函数
def classify_text(text):
model.eval() # 切换到评估模式
inputs = preprocess_text(text)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = torch.softmax(logits, dim=1)
predicted_label = torch.argmax(probabilities, dim=1).item()
return predicted_label, probabilities
# 示例文本
text = "宵凉百念集孤灯,暗雨鸣廊睡未能。生计坐怜秋一叶,归程冥想浪千层。寒心国事浑难料,堆眼官资信可憎。此去梦中应不忘,顺承门内近觚棱。"
# 调用分类函数
predicted_label, probabilities = classify_text(text)
# 输出结果
print(f"预测标签: {predicted_label}")
print(f"概率分布: {probabilities}")
label_0代表唐诗风格 label_1代表宋诗风格
- Downloads last month
- 11
This model does not have enough activity to be deployed to Inference API (serverless) yet. Increase its social
visibility and check back later, or deploy to Inference Endpoints (dedicated)
instead.