Spaces:
Running
Running
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() |