attention-sentiment-classifier / inference_example.py
ericwei's picture
Initial model upload
f3ff1d7 verified
import torch
from transformers import AutoTokenizer
from models.huggingface_model import SentimentClassifierForHuggingFace
# Load the model and tokenizer
model = SentimentClassifierForHuggingFace.from_pretrained("./")
tokenizer = AutoTokenizer.from_pretrained("./")
# Prepare text input
text = "I absolutely loved this movie! The acting was superb."
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
# Run inference
model.eval()
with torch.no_grad():
outputs = model(inputs["input_ids"], return_attention=True, return_dict=True)
# Process results
logits = outputs["logits"]
attention_weights = outputs["attention_weights"]
# Get prediction and confidence
probs = torch.nn.functional.softmax(logits, dim=1)
prediction = torch.argmax(probs, dim=1).item()
confidence = probs[0][prediction].item()
sentiment = "Positive" if prediction == 1 else "Negative"
print(f"Text: {text}")
print(f"Sentiment: {sentiment}")
print(f"Confidence: {confidence:.4f}")
# To visualize attention weights, add matplotlib and seaborn imports
# and use attention_weights to create a heatmap