Update app.py

app.py

@@ -26,6 +26,43 @@ model.eval()
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
 # Функция для предсказания
+
+
+label_mapping_type = {'варианты доставки': 0,
+ 'варианты оплаты': 1,
+ 'возврат средств': 2,
+ 'восстановление пароля': 3,
+ 'время доставки': 4,
+ 'выбор адреса доставки': 5,
+ 'жалоба': 6,
+ 'изменение адреса доставки': 7,
+ 'изменение заказа': 8,
+ 'отзыв': 9,
+ 'отмена заказа': 10,
+ 'отслеживание возврата средств': 11,
+ 'отслеживание заказа': 12,
+ 'подписка на новостную рассылку': 13,
+ 'политика возврата': 14,
+ 'получение информации': 15,
+ 'проблемы с оплатой': 16,
+ 'проблемы с регистрацией': 17,
+ 'проверка платы за отмену': 18,
+ 'проверка счета': 19,
+ 'проверка счетов': 20,
+ 'размещение заказа': 21,
+ 'редактирование учетной записи': 22,
+ 'связь с человеком': 23,
+ 'связь со службой поддержки': 24,
+ 'смена учетной записи': 25,
+ 'создание учетной записи': 26,
+ 'удаление аккаунта': 27}
+
+label_mapping_priority = {'высокий': 0, 'низкий': 1, 'средний': 2}
+
+def get_key_by_value(dictionary, value):
+    reverse_dict = {v: k for k, v in dictionary.items()}
+    return reverse_dict.get(value)
+
 def predict(text):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     input_ids = inputs["input_ids"]
@@ -38,15 +75,14 @@ def predict(text):
     pred_task1 = torch.argmax(logits_task1, dim=1).item()
     pred_task2 = torch.argmax(logits_task2, dim=1).item()
     
-    return {"Task 1 Prediction": pred_task1, "Task 2 Prediction": pred_task2}
+    return {"Тема": get_key_by_value(label_mapping_type, pred_task1), "Приоритет": get_key_by_value(label_mapping_priority, pred_task2)}
 
-# Создание интерфейса с Gradio
 iface = gr.Interface(
     fn=predict, 
-    inputs=gr.Textbox(lines=2, placeholder="Введите текст для анализа..."),  # Обновлено на gr.Textbox
+    inputs=gr.Textbox(lines=2, placeholder="Введите запрос для анализа..."),  # Обновлено на gr.Textbox
     outputs=gr.JSON(),  # Обновлено на gr.JSON
-    title="Multi-Task BERT Model",
-    description="Модель BERT для одновременного решения двух задач: тональность текста и тема.",
+    title="Классификация запроса",
+    description="'варианты доставки', 'варианты оплаты', 'возврат средств', 'восстановление пароля', 'время доставки', 'выбор адреса доставки', 'жалоба', 'изменение адреса доставки', 'изменение заказа', 'отзыв', 'отмена заказа', 'отслеживание возврата средств', 'отслеживание заказа', 'подписка на новостную рассылку', 'политика возврата', 'получение информации', 'проблемы с оплатой', 'проблемы с регистрацией', 'проверка платы за отмену', 'проверка счета', 'размещение заказа', 'редактирование учетной записи', 'связь с человеком', 'связь со службой поддержки', 'смена учетной записи', 'создание учетной записи', 'удаление аккаунта'",
 )
 
 iface.launch()