IMDBChatbot / app.py
sabitizen's picture
Update app.py (#1)
f8fa3b8
Raw
History Blame Contribute Delete
1.28 kB
import torch
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
MODEL_PATH = "sabitizen/distilbert-imdb-movie-review"
# Load model once
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()
def analyze_review(review):
if review.strip() == "":
return "Please enter a movie review."
inputs = tokenizer(
review,
return_tensors="pt",
truncation=True,
padding=True,
max_length=256
)
with torch.no_grad():
outputs = model(**inputs)
probs = F.softmax(outputs.logits, dim=1)
confidence, prediction = torch.max(probs, dim=1)
sentiment = "Positive 😊" if prediction.item() == 1 else "Negative 😞"
return f"""
🎬 **Sentiment:** {sentiment}
πŸ“Š **Confidence:** {confidence.item():.2f}
"""
# Gradio UI
interface = gr.Interface(
fn=analyze_review,
inputs=gr.Textbox(
lines=4,
placeholder="Write a movie review here..."
),
outputs="markdown",
title="🎬 Movie Review Chatbot",
description="DistilBERT fine-tuned on IMDB movie reviews"
)
if __name__ == "__main__":
interface.launch()