Update app.py

app.py

@@ -1,41 +1,39 @@
 import gradio as gr
-from transformers import pipeline
-import re
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+from sklearn.metrics.pairwise import cosine_similarity
 
 MODELS = {
-    "ruRoberta-large": "sberbank-ai/ruRoberta-large",
     "rubert-tiny2": "cointegrated/rubert-tiny2",
-    "multilingual-e5": "intfloat/multilingual-e5-base"
+    "sbert": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+    "LaBSE": "sentence-transformers/LaBSE",
+    "ruRoberta": "sberbank-ai/ruRoberta-large"
 }
 
+def get_embeddings(model, tokenizer, texts):
+    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
+    outputs = model(**inputs)
+    return outputs.last_hidden_state[:, 0].detach().numpy()
+
 def classify(model_name: str, item: str, categories: str) -> str:
-    # Нормализация текста
-    item = re.sub(r"[^а-яА-ЯёЁ]", " ", item).lower().strip()
+    # Загрузка модели и токенизатора
+    tokenizer = AutoTokenizer.from_pretrained(MODELS[model_name])
+    model = AutoModel.from_pretrained(MODELS[model_name])
+    
+    # Подготовка текстов
+    texts = [item] + [c.strip() for c in categories.split(",")]
     
-    classifier = pipeline(
-        "zero-shot-classification",
-        model=MODELS[model_name],
-        device=-1
-    )
+    # Получение эмбеддингов
+    embeddings = get_embeddings(model, tokenizer, texts)
     
-    hypothesis_template = (
-        "Примеры категорий:\n"
-        "- молоток → инструменты\n"
-        "- картофель → овощи\n"
-        "Категория для '{}' → "
-    )
+    # Сравнение с категориями
+    item_embedding = embeddings[0].reshape(1, -1)
+    category_embeddings = embeddings[1:]
     
-    result = classifier(
-        item,
-        candidate_labels=[c.strip().lower() for c in categories.split(",")],
-        hypothesis_template=hypothesis_template,
-        multi_label=False
-    )
+    similarities = cosine_similarity(item_embedding, category_embeddings)[0]
+    best_idx = np.argmax(similarities)
     
-    if result['scores'][0] < 0.3:
-        return "Категория не определена"
-        
-    return f"{result['labels'][0].capitalize()} ({result['scores'][0]:.2f})"
+    return f"{texts[1:][best_idx]} ({similarities[best_idx]:.2f})"
 
 iface = gr.Interface(
     fn=classify,
@@ -44,11 +42,7 @@ iface = gr.Interface(
         gr.Textbox(label="Товар"), 
         gr.Textbox(label="Категории", value="Инструменты, Овощи, Техника")
     ],
-    outputs=gr.Textbox(label="Результат"),
-    examples=[
-        ["ruRoberta-large", "Аккумуляторная дрель", "Инструменты, Техника"],
-        ["rubert-tiny2", "Свёкла кормовая", "Овощи, Фураж"]
-    ]
+    outputs=gr.Textbox(label="Результат")
 )
- 
+
 iface.launch()