import gradio as gr
from transformers import pipeline

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

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

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

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

iface.launch()