Update app.py

app.py

@@ -5,13 +5,15 @@ from transformers import pipeline
 classifier = pipeline("text-classification", model="cointegrated/rubert-tiny2")
 
 # Категории в порядке, соответствующем меткам модели
-CATEGORIES = ["Овощи", "Инструменты", "Фрукты"]  # Измените на ваши категории!
+CATEGORIES = ["Овощи", "Инструменты", "Коробки"]  # Измените на ваши категории!
 
 def classify(item: str, categories: list) -> str:
-    prompt = f"Отнеси '{item}' к категории: {', '.join(categories)}. Ответь только категорией."
+    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]
 
@@ -19,7 +21,7 @@ def classify(item: str, categories: list) -> str:
 iface = gr.Interface(
     fn=classify,
     inputs=[
-        gr.Textbox(label="Название товара"), 
+        gr.Textbox(label="Название товара 1"), 
         gr.Textbox(label="Категории (через запятую)", value="Овощи, Инструменты, Коробки")
     ],
     outputs=gr.Textbox(label="Категория")