import gradio as gr
from transformers import pipeline

# Загружаем модель
classifier = pipeline("text-classification", model="sileod/rubert-tiny-for-ner")

# Категории в порядке, соответствующем меткам модели
CATEGORIES = ["Техника", "Овощи", "Инструменты", "Коробки", "Прочее"]  # Измените на ваши категории!

def classify(item: str, categories: list) -> str:
    prompt = f"""
    Товар: {item}
    Категории: {", ".join(categories)}.
    К какой категории относится товар? Ответь только названием категории.
    """
    result = classifier(prompt, truncation=True)
    label_index = int(result[0]['label'].split("_")[1])
    return CATEGORIES[label_index]

# Интерфейс Gradio
iface = gr.Interface(
    fn=classify,
    inputs=[
        gr.Textbox(label="Название товара rubert-tiny-for-ner"), 
        gr.Textbox(label="Категории (через запятую)", value="Техника, Овощи, Инструменты, Коробки, Прочее")
    ],
    outputs=gr.Textbox(label="Категория")
)

iface.launch()