import gradio as gr
from transformers import pipeline
import re

classifier = pipeline(
    "zero-shot-classification",
    model="cointegrated/rubert-tiny2-2ch",  # Специальная версия для классификации
    device=-1
)

def preprocess(text: str) -> str:
    """Нормализация текста для улучшения качества классификации"""
    text = re.sub(r"[^а-яА-ЯёЁa-zA-Z0-9]", " ", text)  # Удаляем спецсимволы
    text = re.sub(r"\s+", " ", text).strip().lower()    # Нормализуем пробелы и регистр
    return text

def classify(item: str, categories: str) -> str:
    # Предобработка входящих данных
    item = preprocess(item)
    categories_list = [preprocess(c) for c in categories.split(",")]
    
    # Формируем контекстные примеры для улучшения понимания
    hypothesis_template = (
        "Примеры категоризации:\n"
        "- 'молоток' → инструменты\n"
        "- 'морковь' → овощи\n"
        "Теперь определи категорию для: '{}' → "
    )
    
    result = classifier(
        item,
        candidate_labels=categories_list,
        hypothesis_template=hypothesis_template,
        multi_label=False
    )
    
    # Фильтр низкой уверенности
    if result['scores'][0] < 0.5:
        return "Не удалось определить категорию"
        
    return f"{result['labels'][0].capitalize()} (точность: {result['scores'][0]:.2f})"

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

iface.launch()