Update app.py

app.py

@@ -10,10 +10,14 @@ MODELS = {
     "ruRoberta": "sberbank-ai/ruRoberta-large"
 }
 
+PROMPT_TEMPLATES = {
+    "basic": "Товар: {item}. Категория:",
+    "examples": "Примеры:\n- Молоток → Инструменты\n- Морковь → Овощи\nТовар: {item} → ",
+    "strict": "Выбери категорию из [{categories}]. Товар: {item}. Категория:"
+}
+
 def get_embeddings(model, tokenizer, text):
-    # Добавляем промпт
-    prompted_text = f"Товар: {text}. Категория:"
-    inputs = tokenizer(prompted_text, 
+    inputs = tokenizer(text, 
                       padding=True, 
                       truncation=True, 
                       return_tensors="pt",
@@ -21,18 +25,22 @@ def get_embeddings(model, tokenizer, text):
     outputs = model(**inputs)
     return outputs.last_hidden_state[:, 0].detach().numpy()
 
-def classify(model_name: str, item: str, categories: str) -> str:
+def classify(model_name: str, prompt_type: str, item: str, categories: str) -> str:
     tokenizer = AutoTokenizer.from_pretrained(MODELS[model_name])
     model = AutoModel.from_pretrained(MODELS[model_name])
     
-    # Эмбеддинги для товара с промптом
-    item_embedding = get_embeddings(model, tokenizer, item)
+    # Формируем промпт
+    prompt = PROMPT_TEMPLATES[prompt_type].format(
+        item=item,
+        categories=", ".join([c.strip() for c in categories.split(",")])
+    )
     
-    # Эмбеддинги для категорий
-    category_embeddings = []
-    for category in categories.split(","):
-        emb = get_embeddings(model, tokenizer, category.strip())
-        category_embeddings.append(emb)
+    # Эмбеддинги
+    item_embedding = get_embeddings(model, tokenizer, prompt)
+    category_embeddings = [
+        get_embeddings(model, tokenizer, c.strip()) 
+        for c in categories.split(",")
+    ]
     
     # Сравнение
     similarities = cosine_similarity(item_embedding, np.vstack(category_embeddings))[0]
@@ -43,9 +51,10 @@ def classify(model_name: str, item: str, categories: str) -> str:
 gr.Interface(
     fn=classify,
     inputs=[
-        gr.Dropdown(list(MODELS.keys())),
-        gr.Textbox(),
-        gr.Textbox(value="Инструменты, Овощи, Техника")
+        gr.Dropdown(list(MODELS.keys()), label="Модель"),
+        gr.Dropdown(list(PROMPT_TEMPLATES.keys()), label="Шаблон промпта"),
+        gr.Textbox(label="Товар"),
+        gr.Textbox(label="Категории", value="Инструменты, Овощи, Техника")
     ],
     outputs=gr.Textbox()
 ).launch()