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()