Create predict_bert.py

predict_bert.py

@@ -0,0 +1,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()