Spaces:
Sleeping
Sleeping
| 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() | |