import gradio as gr
from transformers import pipeline

MODEL_ID = "arielb30/imdb-distilbert-finetuned"
clf = pipeline("text-classification", model=MODEL_ID)

def predict_sentiment(text):
    if not text.strip():
        return "Please enter some text.", None

    result = clf(text)[0]
    label = result["label"]
    score = float(result["score"])

    pretty = f"Prediction: {label}\nConfidence: {score:.4f}"
    return pretty, {label: score}

demo = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=6, label="Review"),
    outputs=[
        gr.Textbox(label="Result"),
        gr.Label(label="Confidence")
    ],
    title="IMDb Sentiment Classifier",
    description="Paste a movie review and get the model prediction."
)

demo.launch()