from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("moraix/SentimentBot")
model = AutoModelForSequenceClassification.from_pretrained("moraix/SentimentBot").to("cuda")

# Example input
text = "I love this movie so much!"
inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128).to("cuda")

# Predict sentiment
model.eval()
with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.softmax(outputs.logits, dim=1)
    predicted_class = torch.argmax(predictions, dim=1).item()
    confidence = predictions[0, predicted_class].item()

label_map = {0: "negative", 1: "positive"}
sentiment = label_map[predicted_class]
print(f"Text: {text}")
print(f"Sentiment: {sentiment} (Confidence: {confidence:.2f})")