Spaces:
Runtime error
Runtime error
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.")) |