from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer from the Hub
model_name = "FlukeTJ/wangchanberta-base-att-spm-uncased-finetuned-sentiment-cleaned-40k"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Set device (GPU if available, else CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

def predict_sentiment(text):
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # Make prediction
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get probabilities
    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
    # Get the predicted class 
    predicted_class = torch.argmax(probabilities, dim=1).item()
    
    # Map class to sentiment
    sentiment_map = {0: "Neutral", 1: "Positive", 2: "Negative"}
    predicted_sentiment = sentiment_map[predicted_class]
    
    # Get the confidence score
    confidence = probabilities[0][predicted_class].item()
    
    return predicted_sentiment, confidence

# Example usage
texts = [
    "สุดยอดดด"
]

for text in texts:
    sentiment, confidence = predict_sentiment(text)
    print(f"Text: {text}")
    print(f"Predicted Sentiment: {sentiment}")
    print(f"Confidence: {confidence:.2f}")
    print()