File size: 3,395 Bytes
af9853e |
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 83 84 |
import torch
import os
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from .config import Config
class SentimentPredictor:
def __init__(self, model_path=None):
# 1. 如果未指定路径，尝试自动寻找最新的模型
if model_path is None:
# 优先检查 Config.CHECKPOINT_DIR (如果训练完成，final_model 会在这里)
if os.path.exists(os.path.join(Config.CHECKPOINT_DIR, "config.json")):
model_path = Config.CHECKPOINT_DIR
else:
# 如果没有 final_model，尝试寻找最新的 checkpoint (在 results 目录)
import glob
ckpt_list = glob.glob(os.path.join(Config.RESULTS_DIR, "checkpoint-*"))
if ckpt_list:
# 按修改时间排序，取最新的
ckpt_list.sort(key=os.path.getmtime)
model_path = ckpt_list[-1]
print(f"Using latest checkpoint found: {model_path}")
else:
# 只有在真的找不到时才回退
model_path = Config.CHECKPOINT_DIR
print(f"Loading model from {model_path}...")
try:
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
except OSError:
print(f"Warning: Model not found at {model_path}. Loading base model for demo purpose.")
self.tokenizer = AutoTokenizer.from_pretrained(Config.BASE_MODEL)
self.model = AutoModelForSequenceClassification.from_pretrained(Config.BASE_MODEL, num_labels=Config.NUM_LABELS)
# Device selection
if torch.backends.mps.is_available():
self.device = torch.device("mps")
elif torch.cuda.is_available():
self.device = torch.device("cuda")
else:
self.device = torch.device("cpu")
self.model.to(self.device)
self.model.eval()
def predict(self, text):
inputs = self.tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=Config.MAX_LENGTH,
padding=True
)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = self.model(**inputs)
probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
prediction = torch.argmax(probabilities, dim=-1).item()
score = probabilities[0][prediction].item()
label = Config.ID2LABEL.get(prediction, "unknown")
return {
"text": text,
"sentiment": label,
"confidence": f"{score:.4f}"
}
if __name__ == "__main__":
# Demo
predictor = SentimentPredictor()
test_texts = [
"这家店的快递太慢了，而且东西味道很奇怪。",
"非常不错，包装很精美，下次还会来买。",
"感觉一般般吧，没有想象中那么好，但也还可以。"
]
print("\nPredicting...")
for text in test_texts:
result = predictor.predict(text)
print(f"Text: {result['text']}")
print(f"Sentiment: {result['sentiment']} (Confidence: {result['confidence']})")
print("-" * 30)
|