Yu,Zixuan

Create predict_bert.py

0ee5b66 12 months ago

2.33 kB

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

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