import gradio as gr
from transformers import pipeline
import time

# Загружаем модель ПЕРВОЙ
try:
    classifier = pipeline(
        "text-classification", 
        model="cointegrated/rubert-tiny2",
        truncation=True
    )
except Exception as e:
    print(f"Ошибка загрузки модели: {e}")
    raise

# Функция классификации
def classify(item: str, categories: str) -> str:
    try:
        prompt = f"Отнеси '{item}' к категории: {categories}. Ответь только категорией."
        result = classifier(prompt)
        return result[0]['label']
    except Exception as e:
        return f"Ошибка: {str(e)}"

# Создаем интерфейс ПОСЛЕ загрузки модели
iface = gr.Interface(
    fn=classify,
    inputs=[
        gr.Textbox(label="Товар"), 
        gr.Textbox(label="Категории (через запятую)", value="Овощи, Инструменты")
    ],
    outputs="text",
    title="Классификатор"
)

# Явно запускаем сервер
iface.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True
)

# Бесконечный цикл, чтобы сервер не завершался
while True:
    time.sleep(60)