NomClass / app.py
hh1199's picture
Update app.py
19a9f33 verified
raw
history blame
675 Bytes
import gradio as gr
from transformers import pipeline
classifier = pipeline("zero-shot-classification", model="cointegrated/rubert-tiny2")
def classify(item: str, categories: str) -> str:
categories_list = [c.strip() for c in categories.split(",")]
result = classifier(item, categories_list, multi_label=False)
return result['labels'][0]
iface = gr.Interface(
fn=classify,
inputs=[
gr.Textbox(label="Название товара"),
gr.Textbox(label="Категории (через запятую)", value="Техника, Овощи, Инструменты")
],
outputs=gr.Textbox(label="Категория")
)
iface.launch()