import gradio as gr
from transformers import pipeline

model = pipeline(task="sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")

def sentiment_analysis(text):
    res = model(text)[0]
    res_label = {}
    if res["label"] == "POSITIVE":
        res_label["POSITIVE"] = res["score"]
        res_label["NEGATIVE"] = 1 - res["score"]
    if res["label"] == "NEGATIVE":
        res_label["NEGATIVE"] = res["score"]
        res_label["POSITIVE"] = 1 - res["score"]
    return res_label

custom_css = """
#component-0 {
    max-width: 600px;
    margin: 0 auto;
}

h1,h2 {
    text-align: center;
}

a {
    color: #77b3ee !important;
    text-decoration: none !important;
}

a:hover {
    text-decoration: underline !important;
}
"""

browser_tab_title = "Sentiment Analysis"
intro_markdown = """## Sentiment Analysis

Using the [distilbert-base-uncased-finetuned-sst-2-english](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english) model, trained on movie reviews."""

with gr.Blocks(title=browser_tab_title, css=custom_css) as demo:
    with gr.Row():
        with gr.Column():
            title = gr.Markdown(intro_markdown)
            text_input = gr.Textbox(placeholder="Enter a positive or negative sentence here...", label="Text")
            label_output = gr.Label(label="Sentiment outcome")
            button_run = gr.Button("Compute sentiment")
            button_run.click(sentiment_analysis, inputs=text_input, outputs=label_output)
            gr.Examples(["That's great!", "The movie was bad.", "How are you"], text_input)

if __name__ == "__main__":
    demo.launch()