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