from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Load the trained model
tokenizer = AutoTokenizer.from_pretrained("Kunalatmosoft/imdb_model")
model = AutoModelForSequenceClassification.from_pretrained("Kunalatmosoft/imdb_model")

# Function to predict sentiment
def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=1).item()
    return "Positive" if prediction == 1 else "Negative"

# Example usage
print(predict_sentiment("This movie was amazing!"))  
print(predict_sentiment("I didn't like this movie."))