import gradio as gr
from transformers import pipeline

classifier = pipeline(
    "text-classification", model="rasyosef/roberta-base-finetuned-sst2"
)


def predict_sentiment(text):
    return classifier(text)[0]


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Sentiment Classifier

        This model is a fine-tuned version of roberta-base on 
        the glue sst2 dataset for sentiment classification. 
        The model classifies the input text as having either 
        `positive` or `negative` sentiment. 
        """
    )
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="Input text", placeholder="Enter text here", lines=3)
            btn = gr.Button("Classify")
        with gr.Column():
            out = gr.Textbox(label="Sentiment")

    btn.click(fn=predict_sentiment, inputs=inp, outputs=out)
    gr.Examples(
        examples=["This movie was awesome.", "The movie was boring."],
        inputs=[inp],
    )

demo.launch()