NLP_IMDB / app.py
meetmendapara's picture
Refactor sentiment prediction function to improve clarity and confidence extraction
6e547ba
# This script creates a Gradio web app for sentiment analysis of movie reviews using a pre-trained BERT model.
# Import necessary libraries
from transformers import BertTokenizer, BertForSequenceClassification
import torch
import gradio as gr
# Load saved model and tokenizer
model = BertForSequenceClassification.from_pretrained("./imdb_bert_model")
tokenizer = BertTokenizer.from_pretrained("./imdb_bert_model")
# Prediction function
def predict_sentiment(text):
"""
Predicts the sentiment of the given text using the fine-tuned BERT model.
Args:
text (str): The input movie review text.
Returns:
str: The predicted sentiment with confidence.
"""
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
outputs = model(**inputs)
logits = outputs.logits
# Extract the maximum value (confidence) and its index (prediction)
confidence, prediction = torch.max(logits, dim=1)
confidence = confidence.item() # Convert tensor to Python float
prediction = prediction.item() # Convert tensor to Python int
# confidence = torch.max(logits, dim=1).item()
# prediction = torch.argmax(logits, dim=1).item()
sentiment = "Positive 😊" if prediction == 1 else "Negative 😠"
return f"{sentiment} with confidence {confidence * 100:.2f}% confidence"
# Responsive UI with gr.Blocks
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("## 🎬 IMDB Movie Review Sentiment Analyzer")
gr.Markdown("Write a movie review below and let BERT tell you if it's **Positive** or **Negative** 🎯")
with gr.Row():
with gr.Column(scale=2):
review_input = gr.Textbox(
label="Enter Review",
placeholder="e.g. This movie had me on the edge of my seat!",
lines=5,
max_lines=8,
autofocus=True
)
submit_btn = gr.Button("πŸ” Analyze")
with gr.Column(scale=1):
result_output = gr.Label(label="Predicted Sentiment")
gr.Examples(
examples=[
["This movie was absolutely amazing and so emotional!"],
["Worst film I’ve ever seen. Total waste of time."],
["The story was okay, but the acting saved it."],
["A beautiful piece of storytelling. I loved it!"],
],
inputs=[review_input]
)
submit_btn.click(fn=predict_sentiment, inputs=review_input, outputs=result_output)
gr.Markdown("### Made with ❀️ by [Meet Mendapara](https://github.com/Meetmendapara09)")
demo.launch(share=True)