Views
No views yet
1from transformers import BertTokenizer, BertForSequenceClassification
2import torch
3
4# 加载已训练的模型和分词器
5model_path = 'qixun/tangsong_poem_classify'
6tokenizer = BertTokenizer.from_pretrained(model_path)
7model = BertForSequenceClassification.from_pretrained(model_path)
8
9# 预处理函数
10def preprocess_text(text):
11 inputs = tokenizer(text, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
12 return inputs
13
14# 分类函数
15def classify_text(text):
16 model.eval() # 切换到评估模式
17 inputs = preprocess_text(text)
18 with torch.no_grad():
19 outputs = model(**inputs)
20 logits = outputs.logits
21 probabilities = torch.softmax(logits, dim=1)
22 predicted_label = torch.argmax(probabilities, dim=1).item()
23 return predicted_label, probabilities
24
25# 示例文本
26text = "宵凉百念集孤灯，暗雨鸣廊睡未能。生计坐怜秋一叶，归程冥想浪千层。寒心国事浑难料，堆眼官资信可憎。此去梦中应不忘，顺承门内近觚棱。"
27
28# 调用分类函数
29predicted_label, probabilities = classify_text(text)
30
31# 输出结果
32print(f"预测标签: {predicted_label}")
33print(f"概率分布: {probabilities}")
34