NomClass / app.py
hh1199's picture
Update app.py
3ee972d verified
raw
history blame
668 Bytes
import gradio as gr
from transformers import pipeline
models = {
"ruBert-tiny2": "cointegrated/rubert-tiny2",
"ruRoberta-large": "sberbank-ai/ruRoberta-large",
"multilingual-e5": "intfloat/multilingual-e5-base"
}
def classify(model, text, labels):
classifier = pipeline("zero-shot-classification", model=models[model])
result = classifier(text, [l.strip() for l in labels.split(",")])
return result['labels'][0]
gr.Interface(
fn=classify,
inputs=[
gr.Dropdown(list(models.keys())),
gr.Textbox(),
gr.Textbox(value="Овощи, Инструменты, Техника")
],
outputs=gr.Textbox()
).launch()