File size: 668 Bytes
7ab3fbc
 
 
3ee972d
 
 
 
 
7ab3fbc
3ee972d
 
 
 
1515adb
3ee972d
7ab3fbc
 
3ee972d
 
 
86bd5a4
3ee972d
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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()