qixun's picture
Update README.md
ad38916 verified
metadata
license: gpl-3.0
widget:
  - text: 宵凉百念集孤灯,暗雨鸣廊睡未能。生计坐怜秋一叶,归程冥想浪千层。寒心国事浑难料,堆眼官资信可憎。此去梦中应不忘,顺承门内近觚棱。
pipeline_tag: text-classification

使用方法如下:

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代表宋诗风格