import gradio as gr

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

tokenizer = AutoTokenizer.from_pretrained("ua-l/topics-classifier")
model = AutoModelForSequenceClassification.from_pretrained(
    "ua-l/topics-classifier"
)

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 str(topics)

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