File size: 2,179 Bytes
fba3401 | 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 82 | import torch
from transformers import BertTokenizer
from model.sentiment_model import SentimentAnalysisModel
# 设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. 加载 tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
# 2. 加载模型结构
model = SentimentAnalysisModel("bert-base-uncased")
# 3. 加载训练好的权重
model.load_state_dict(
torch.load("bert_imdb_sentiment.pth", map_location=device)
)
model.to(device)
model.eval() # ⚠️ 非常重要
print("Model loaded successfully.")
def predict_sentiment(text):
inputs = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=256,
return_tensors="pt"
)
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
with torch.no_grad():
outputs = model(input_ids, attention_mask)
probs = torch.softmax(outputs, dim=1)
pred = torch.argmax(probs, dim=1).item()
label_map = {0: "Negative 😡", 1: "Positive 😊"}
return label_map[pred], probs[0][pred].item()
def batch_predict(texts):
inputs = tokenizer(
texts,
padding=True,
truncation=True,
max_length=256,
return_tensors="pt"
)
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
with torch.no_grad():
outputs = model(input_ids, attention_mask)
preds = torch.argmax(outputs, dim=1)
return preds.cpu().tolist()
if __name__ == "__main__":
# text = "This movie was absolutely amazing, I loved it!"
texts = [
"This movie is terrible.",
"I really enjoyed this film!",
"Not bad, but could be better."
]
results = batch_predict(texts)
print(results) # [0, 1, 1]
# | 数值 | 含义 |
# | -- | -------- |
# | 0 | Negative |
# | 1 | Positive |
# label, confidence = predict_sentiment(texts)
# print(f"Text: {text}")
# print(f"Prediction: {label}, confidence: {confidence:.4f}")
|