Update app.py

app.py

@@ -1,23 +1,54 @@
 import gradio as gr
 from transformers import pipeline
+import re
 
-models = {
-    "ruBert-tiny2": "cointegrated/rubert-tiny2",
+MODELS = {
     "ruRoberta-large": "sberbank-ai/ruRoberta-large",
+    "rubert-tiny2": "cointegrated/rubert-tiny2",
     "multilingual-e5": "intfloat/multilingual-e5-base"
 }
 
-def classify(model, text, labels):
-    classifier = pipeline("zero-shot-classification", model=models[model])
-    result = classifier(text, [l.strip() for l in labels.split(",")])
-    return result['labels'][0]
+def classify(model_name: str, item: str, categories: str) -> str:
+    # Нормализация текста
+    item = re.sub(r"[^а-яА-ЯёЁ]", " ", item).lower().strip()
+    
+    classifier = pipeline(
+        "zero-shot-classification",
+        model=MODELS[model_name],
+        device=-1
+    )
+    
+    hypothesis_template = (
+        "Примеры категорий:\n"
+        "- молоток → инструменты\n"
+        "- картофель → овощи\n"
+        "Категория для '{}' → "
+    )
+    
+    result = classifier(
+        item,
+        candidate_labels=[c.strip().lower() for c in categories.split(",")],
+        hypothesis_template=hypothesis_template,
+        multi_label=False
+    )
+    
+    if result['scores'][0] < 0.3:
+        return "Категория не определена"
+        
+    return f"{result['labels'][0].capitalize()} ({result['scores'][0]:.2f})"
 
-gr.Interface(
+iface = gr.Interface(
     fn=classify,
     inputs=[
-        gr.Dropdown(list(models.keys())),
-        gr.Textbox(),
-        gr.Textbox(value="Овощи, Инструменты, Техника")
+        gr.Dropdown(list(MODELS.keys()), label="Модель"),
+        gr.Textbox(label="Товар"), 
+        gr.Textbox(label="Категории", value="Инструменты, Овощи, Техника")
     ],
-    outputs=gr.Textbox()
-).launch()
+    outputs=gr.Textbox(label="Результат"),
+    examples=[
+        ["ruRoberta-large", "Аккумуляторная дрель", "Инструменты, Техника"],
+        ["rubert-tiny2", "Свёкла кормовая", "Овощи, Фураж"]
+    ]
+)
+
+iface.launch()