import gradio as gr
from transformers import pipeline

classifier = pipeline(
    "zero-shot-classification",
    model="cointegrated/rubert-tiny2",
    device=-1  # Используем CPU для стабильности на бесплатных инстансах
)

def classify(item: str, categories: str) -> str:
    # Формируем гипотезу для классификации
    hypothesis_template = "Этот текст относится к категории {}."
    
    categories_list = [c.strip() for c in categories.split(",")]
    
    result = classifier(
        item,
        candidate_labels=categories_list,
        hypothesis_template=hypothesis_template,
        multi_label=False
    )
    
    # Возвращаем категорию с наибольшим скором
    return f"{result['labels'][0]} (уверенность: {result['scores'][0]:.2f})"

iface = gr.Interface(
    fn=classify,
    inputs=[
        gr.Textbox(label="Название товара", placeholder="Введите наименование товара"), 
        gr.Textbox(label="Категории (через запятую)", 
                 value="Овощи, Инструменты, Техника, Упаковка")
    ],
    outputs=gr.Textbox(label="Результат классификации"),
    examples=[
        ["Молоток", "Овощи, Инструменты, Техника"],
        ["Свекла", "Фрукты, Овощи, Ягоды"],
        ["Ноутбук", "Техника, Мебель, Инструменты"]
    ],
    title="Классификатор товаров",
    description="Введите название товара и список категорий через запятую"
)

iface.launch(debug=True)