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