Update app.py

app.py

@@ -10,39 +10,42 @@ MODELS = {
     "ruRoberta": "sberbank-ai/ruRoberta-large"
 }
 
-def get_embeddings(model, tokenizer, texts):
-    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
+def get_embeddings(model, tokenizer, text):
+    # Добавляем промпт
+    prompted_text = f"Товар: {text}. Категория:"
+    inputs = tokenizer(prompted_text, 
+                      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:
-    # Загрузка модели и токенизатора
     tokenizer = AutoTokenizer.from_pretrained(MODELS[model_name])
     model = AutoModel.from_pretrained(MODELS[model_name])
     
-    # Подготовка текстов
-    texts = [item] + [c.strip() for c in categories.split(",")]
+    # Эмбеддинги для товара с промптом
+    item_embedding = get_embeddings(model, tokenizer, item)
     
-    # Получение эмбеддингов
-    embeddings = get_embeddings(model, tokenizer, texts)
+    # Эмбеддинги для категорий
+    category_embeddings = []
+    for category in categories.split(","):
+        emb = get_embeddings(model, tokenizer, category.strip())
+        category_embeddings.append(emb)
     
-    # Сравнение с категориями
-    item_embedding = embeddings[0].reshape(1, -1)
-    category_embeddings = embeddings[1:]
-    
-    similarities = cosine_similarity(item_embedding, category_embeddings)[0]
+    # Сравнение
+    similarities = cosine_similarity(item_embedding, np.vstack(category_embeddings))[0]
     best_idx = np.argmax(similarities)
     
-    return f"{texts[1:][best_idx]} ({similarities[best_idx]:.2f})"
+    return f"{categories.split(',')[best_idx].strip()} ({similarities[best_idx]:.2f})"
 
-iface = gr.Interface(
+gr.Interface(
     fn=classify,
     inputs=[
-        gr.Dropdown(list(MODELS.keys()), label="Модель"),
-        gr.Textbox(label="Товар"), 
-        gr.Textbox(label="Категории", value="Инструменты, Овощи, Техника")
+        gr.Dropdown(list(MODELS.keys())),
+        gr.Textbox(),
+        gr.Textbox(value="Инструменты, Овощи, Техника")
     ],
-    outputs=gr.Textbox(label="Результат")
-)
-
-iface.launch()
+    outputs=gr.Textbox()
+).launch()