# 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)