Update app.py

app.py

@@ -1,30 +1,18 @@
 import gradio as gr
 from transformers import pipeline
 
-# Загружаем модель
-classifier = pipeline("text-classification", model="cointegrated/rubert-tiny2")
+classifier = pipeline("zero-shot-classification", model="cointegrated/rubert-tiny2")
 
-# Категории в порядке, соответствующем меткам модели
-CATEGORIES = ["Техника", "Овощи", "Инструменты", "Коробки", "Прочее"]  # Измените на ваши категории!
+def classify(item: str, categories: str) -> str:
+    categories_list = [c.strip() for c in categories.split(",")]
+    result = classifier(item, categories_list, multi_label=False)
+    return result['labels'][0]
 
-def classify(item: str, categories: list) -> str:
-    prompt = f"""
-    Товар: {item}
-    Категории: {", ".join(categories)}.
-    К какой категории относится товар? Ответь только названием категории.
-    """
-    result = classifier(prompt, truncation=True)
-    print(result) 
-    label_index = int(result[0]['label'].split("_")[1])
-    #return CATEGORIES[label_index]
-    return result
-
-# Интерфейс Gradio
 iface = gr.Interface(
     fn=classify,
     inputs=[
-        gr.Textbox(label="Название товара cointegrated/rubert-tiny2 ++"), 
-        gr.Textbox(label="Категории (через запятую)", value="Техника, Овощи, Инструменты, Коробки, Прочее")
+        gr.Textbox(label="Название товара"), 
+        gr.Textbox(label="Категории (через запятую)", value="Техника, Овощи, Инструменты")
     ],
     outputs=gr.Textbox(label="Категория")
 )