import gradio as gr
from transformers import pipeline

# Load the sentiment analysis pipeline
classifier = pipeline("sentiment-analysis", model="Gholamreza/distillbert-imdb")

def predict_sentiment(text):
    result = classifier(text)
    label = result[0]['label']
    score = result[0]['score']
    print(label, score)
    if label == "LABEL_0":
        score = 1 - score  # Adjust score for negative
    return {"Probability Positive": score, "Probability Negative": 1 - score}

# Define Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## IMDB Sentiment Classifier")

    with gr.Row():
        text_input = gr.Textbox(label="Input Review")

    with gr.Row():
        predict_button = gr.Button("Predict")

    with gr.Row():
        output = gr.Label(label="Sentiment Probabilities")

    examples = [
        "The movie was fantastic! The performances were top-notch and the story was gripping.",
        "I didn't enjoy the movie. The plot was predictable and the acting was subpar.",
        "It was an average film. Some parts were good, but it didn't leave a lasting impression."
    ]

    gr.Examples(examples=examples, inputs=text_input)
    predict_button.click(fn=predict_sentiment, inputs=text_input, outputs=output)

# Launch the app
demo.launch()