import gradio as gr
from transformers import pipeline
from handler import EndpointHandler

handler = EndpointHandler('bie-nhd/visobert-multitask')

def predict(text, task):
    # Call HuggingFace Inference Endpoint to get classification result
    result = handler({"task": task, "text": text})
    return result

with gr.Blocks() as app:
    gr.Markdown("# Text Classification")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="text")
            task_dropdown = gr.Dropdown(label="task", choices=["all", "sentiment", "topic", "hate_speech"])
            predict_btn = gr.Button("Predict")
        with gr.Column():
            output = gr.Textbox(label="Result")

    predict_btn.click(fn=predict, inputs=[text_input, task_dropdown], outputs=output)


app.launch(share=True)