test / app.py
Kunalatmosoft's picture
again update
5cfb632 verified
raw
history blame contribute delete
687 Bytes
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."))