File size: 2,327 Bytes
0ee5b66 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import torch
from transformers import BertTokenizer, BertForSequenceClassification
import re
def preprocess_text(text):
text = re.sub(r'[^\w\s]', '', text)
return text
def load_model():
# 加载预训练模型和分词器
model_name = 'bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)
# 加载训练好的模型参数
model.load_state_dict(torch.load('best_bert_model.pt'))
model.eval()
return model, tokenizer
def predict_text(text, model, tokenizer, device='cpu'):
# 预处理文本
processed_text = preprocess_text(text)
# 编码文本
encoding = tokenizer.encode_plus(
processed_text,
add_special_tokens=True,
max_length=128,
return_token_type_ids=False,
padding='max_length',
truncation=True,
return_attention_mask=True,
return_tensors='pt'
)
# 移动到设备
input_ids = encoding['input_ids'].to(device)
attention_mask = encoding['attention_mask'].to(device)
# 预测
with torch.no_grad():
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask
)
_, preds = torch.max(outputs.logits, dim=1)
probs = torch.softmax(outputs.logits, dim=1)
return preds.item(), probs[0].cpu().numpy()
def main():
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")
# 加载模型
print("正在加载模型...")
model, tokenizer = load_model()
model = model.to(device)
print("模型加载完成！")
while True:
# 获取用户输入
text = input("\n请输入要预测的文本（输入'q'退出）: ")
if text.lower() == 'q':
break
# 预测
pred, probs = predict_text(text, model, tokenizer, device)
# 输出结果
print("\n预测结果：")
print(f"文本: {text}")
print(f"预测标签: {pred}")
print(f"预测概率: 标签0的概率: {probs[0]:.4f}, 标签1的概率: {probs[1]:.4f}")
print("-" * 50)
if __name__ == "__main__":
main() |