import gradio as gr
from transformers import pipeline

# Загрузка модели
classifier = pipeline("text-classification", model="cointegrated/rubert-tiny2")

def classify(item: str, categories: str) -> str:
prompt = f"Отнеси '{item}' к категории: {categories}. Ответь только категорией."
result = classifier(prompt, truncation=True)
return result[0]['label']

# Интерфейс
iface = gr.Interface(
fn=classify,
inputs=[
gr.Textbox(label="Название товара"),
gr.Textbox(label="Категории (через запятую)", value="Овощи, Инструменты, Фрукты")
],
outputs="text",
title="Классификатор номенклатуры",
examples=[["Морковь", "Овощи, Инструменты"], ["Дрель", "Инструменты, Электроника"]]
)

iface.launch()