import gradio as gr
import torch
from transformers import pipeline


pipe = pipeline("text-classification",  # tarea a realizar
            model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",  # modelo a utilizar
            torch_dtype=torch.bfloat16,)   #


def text_classification(text):
    result = pipe(text)
    return result[0]['label']


examples = ['You are the evil!', 'Doing great boy!']

demo = gr.Interface(fn=text_classification,
                    inputs= gr.Textbox(lines=2, label="Text", placeholder="Enter text here..."), 
                    outputs=gr.Textbox(lines=2, label="Text Classification Result"),
                    description="Enter a text and see the sentiment classification result!",
                    examples=examples,
                    title="Sentiment Classification")

demo.launch()