Update app.py

app.py

@@ -3,23 +3,50 @@ import time
 from transformers import pipeline
 from datasets import load_dataset
 
-# Примеры вручную — заменяют шумные выборки из датасета
-few_shot_examples = [
-    "Клиент: Я не могу войти в личный кабинет\nОтвет: Пожалуйста, проверьте правильность логина и пароля. Если проблема сохраняется — воспользуйтесь восстановлением доступа или обратитесь в поддержку.",
-    "Клиент: Мне пришло смс о входе, которого я не совершал\nОтвет: Это может быть попытка входа третьего лица. Рекомендуем срочно сменить пароль и обратиться в службу поддержки банка."
-]
+# Группированные обращения по категориям
+categories = {
+    "доступ": [
+        "Клиент: Я не могу войти в личный кабинет\nОтвет: Пожалуйста, проверьте правильность логина и пароля. Если проблема сохраняется — воспользуйтесь восстановлением доступа или обратитесь в поддержку.",
+        "Клиент: У меня заблокирован вход в интернет-банк\nОтвет: Обратитесь в поддержку для подтверждения личности и восстановления доступа."
+    ],
+    "переводы": [
+        "Клиент: Как перевести деньги на другую карту?\nОтвет: Перевод возможен через мобильное приложение или в отделении. Убедитесь, что карта получателя активна.",
+        "Клиент: Почему не проходит перевод?\nОтвет: Проверьте лимиты по карте и правильность реквизитов. Если всё верно, свяжитесь с поддержкой."
+    ],
+    "смс": [
+        "Клиент: Мне пришло смс, которого я не ожидал\nОтвет: Это может быть техническое уведомление. Уточните дату и текст сообщения, чтобы мы проверили.",
+        "Клиент: Получаю подозрительные смс от банка\nОтвет: Не переходите по ссылкам. Немедленно смените пароль и сообщите в службу безопасности."
+    ]
+}
 
-# Инструкции
+# Инструкция CoT (расширенная)
 cot_instruction = (
-    "Ты — вежливый банковский ассистент. Клиент описывает проблему. "
-    "Твоя задача — сначала шаг за шагом проанализировать суть запроса, а затем дать полезный и точный ответ. "
-    "Если информации недостаточно — запроси уточнение. Не выдумывай факты."
+    "Ты — банковский ассистент. Клиент задал вопрос. Сначала разложи его по смысловым блокам: что именно он хочет, какие данные нужны, возможные причины. "
+    "После анализа предложи чёткий, полезный и вежливый ответ. Если информации недостаточно — запроси уточнение."
 )
 
+# Инструкция обычная
 simple_instruction = (
-    "Ты — банковский помощник. Отвечай кратко и понятно. Не используй лишние детали."
+    "Ты — банковский помощник. Отвечай официально, коротко и понятно, без рассуждений."
 )
 
+# Классификация темы обращения
+
+def classify_topic(user_input):
+    if any(x in user_input.lower() for x in ["войти", "кабинет", "доступ", "логин", "пароль"]):
+        return "доступ"
+    elif any(x in user_input.lower() for x in ["перевод", "перевести", "деньги", "карта"]):
+        return "переводы"
+    elif any(x in user_input.lower() for x in ["смс", "сообщение", "уведомление"]):
+        return "смс"
+    return "доступ"  # fallback
+
+# Сбор примеров по теме
+
+def get_few_shot_examples(user_input):
+    topic = classify_topic(user_input)
+    return categories.get(topic, categories["доступ"])
+
 # Модели
 models = {
     "ruDialoGPT-small": pipeline("text-generation", model="t-bank-ai/ruDialoGPT-small", tokenizer="t-bank-ai/ruDialoGPT-small", device=-1),
@@ -28,37 +55,40 @@ models = {
 }
 
 # Формирование промптов
+
 def build_cot_prompt(user_input):
-    examples = "\n\n".join(few_shot_examples)
+    examples = "\n\n".join(get_few_shot_examples(user_input))
     return f"{cot_instruction}\n\n{examples}\n\nКлиент: {user_input}\nРассуждение и ответ:"
 
 def build_simple_prompt(user_input):
-    examples = "\n\n".join(few_shot_examples)
+    examples = "\n\n".join(get_few_shot_examples(user_input))
     return f"{simple_instruction}\n\n{examples}\n\nКлиент: {user_input}\nОтвет:"
 
-# Функция извлечения финального ответа
+# Фильтрация финального ответа
+
 def extract_final_answer(generated_text):
     for line in generated_text.split('\n'):
         if "Ответ:" in line:
             return line.split("Ответ:")[-1].strip()
     return generated_text.strip().split('\n')[-1]
 
-# Генерация по обоим промптам
+# Генерация
+
 def generate_dual_answers(user_input):
     results = {}
     prompt_cot = build_cot_prompt(user_input)
     prompt_simple = build_simple_prompt(user_input)
 
     for name, pipe in models.items():
-        # CoT ответ
+        # CoT
         start_cot = time.time()
-        out_cot = pipe(prompt_cot, max_new_tokens=80, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        out_cot = pipe(prompt_cot, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.9)[0]["generated_text"]
         end_cot = round(time.time() - start_cot, 2)
         answer_cot = extract_final_answer(out_cot)
 
-        # Обычный ответ
+        # Simple
         start_simple = time.time()
-        out_simple = pipe(prompt_simple, max_new_tokens=80, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        out_simple = pipe(prompt_simple, max_new_tokens=100, do_sample=True, top_p=0.95, temperature=0.9)[0]["generated_text"]
         end_simple = round(time.time() - start_simple, 2)
         answer_simple = extract_final_answer(out_simple)
 
@@ -80,9 +110,9 @@ def generate_dual_answers(user_input):
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏦 Банковский помощник: сравнение CoT и обычного ответа")
+    gr.Markdown("## 🏦 Банковский помощник: CoT vs. Обычный ответ (магистерская работа)")
 
-    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Как перевести деньги на другую карту?", lines=2)
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Почему не проходит перевод?", lines=2)
     btn = gr.Button("Сгенерировать")
 
     gr.Markdown("### ruDialoGPT-small")
@@ -110,3 +140,4 @@ with gr.Blocks() as demo:
     ])
 
 demo.launch()
+