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