import gradio as gr

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

tokenizer = AutoTokenizer.from_pretrained("ukr-models/xlm-roberta-base-uk")
model = AutoModelForSequenceClassification.from_pretrained(
    "ua-l/topics-classifier-xlm-roberta-base-uk-v2"
)

topic_classifier = pipeline(
    task="text-classification", model=model, tokenizer=tokenizer, device="cpu", top_k=5
)

def predict(question):
    predictions = topic_classifier(question)

    topics = []
    for prediction in predictions[0]:
        label = prediction["label"]
        probability = round(prediction["score"] * 100, 2)

        topics.append(
            {
                "label": label,
                "probability": probability,
            }
        )

    return topics

inputs = gr.Textbox(lines=2, label="Enter the text", value="Як отримати виплати ВПО?")

outputs = gr.JSON(label="Output")

demo = gr.Interface(fn=predict, inputs=inputs, outputs=outputs)
demo.launch()