Update app.py

app.py

@@ -3,25 +3,21 @@ import time
 from transformers import pipeline
 from datasets import load_dataset
 
-# Загружаем датасет
-dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
+# Примеры вручную — заменяют шумные выборки из датасета
+few_shot_examples = [
+    "Клиент: Я не могу войти в личный кабинет\nОтвет: Пожалуйста, проверьте правильность логина и пароля. Если проблема сохраняется — воспользуйтесь восстановлением доступа или обратитесь в поддержку.",
+    "Клиент: Мне пришло смс о входе, которого я не совершал\nОтвет: Это может быть попытка входа третьего лица. Рекомендуем срочно сменить пароль и обратиться в службу поддержки банка."
+]
 
-# Примеры для few-shot (без 'rating')
-few_shot_examples = []
-for row in dataset.select(range(2)):
-    review = row["review"]
-    ex = f"Клиент: {review}\nОтвет: Спасибо за обращение! Уточните, пожалуйста, детали ситуации, чтобы мы могли помочь."
-    few_shot_examples.append(ex)
-
-# Системные инструкции
+# Инструкции
 cot_instruction = (
-    "Ты — вежливый банковский помощник. Клиент описывает проблему. "
-    "Сначала проанализируй её шаг за шагом, потом сформулируй итоговый ответ."
+    "Ты — вежливый банковский ассистент. Клиент описывает проблему. "
+    "Твоя задача — сначала шаг за шагом проанализировать суть запроса, а затем дать полезный и точный ответ. "
+    "Если информации недостаточно — запроси уточнение. Не выдумывай факты."
 )
 
 simple_instruction = (
-    "Ты — вежливый банковский помощник. Отвечай кратко и официально, без лишних деталей. "
-    "Не используй рассуждение, просто дай понятный клиенту ответ."
+    "Ты — банковский помощник. Отвечай кратко и понятно. Не используй лишние детали."
 )
 
 # Модели
@@ -31,41 +27,40 @@ models = {
     "ruGPT3-small": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_based_on_gpt2", device=-1),
 }
 
-# Промпт CoT
+# Формирование промптов
 def build_cot_prompt(user_input):
     examples = "\n\n".join(few_shot_examples)
-    return (
-        f"{cot_instruction}\n\n{examples}\n\nКлиент: {user_input}\n"
-        "Рассуждение и ответ:"
-    )
+    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)
-    return (
-        f"{simple_instruction}\n\n{examples}\n\nКлиент: {user_input}\n"
-        "Ответ:"
-    )
+    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=100, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        out_cot = pipe(prompt_cot, max_new_tokens=80, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
         end_cot = round(time.time() - start_cot, 2)
-        answer_cot = out_cot.strip().split('\n')[-1]
+        answer_cot = extract_final_answer(out_cot)
 
-        # Simple
+        # Обычный ответ
         start_simple = time.time()
-        out_simple = pipe(prompt_simple, max_new_tokens=100, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        out_simple = pipe(prompt_simple, max_new_tokens=80, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
         end_simple = round(time.time() - start_simple, 2)
-        answer_simple = out_simple.strip().split('\n')[-1]
+        answer_simple = extract_final_answer(out_simple)
 
         results[name] = {
             "cot_answer": answer_cot,
@@ -85,9 +80,9 @@ def generate_dual_answers(user_input):
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏦 Банковский помощник (2 промпта: рассуждение + краткий ответ)")
+    gr.Markdown("## 🏦 Банковский помощник: сравнение CoT и обычного ответа")
 
-    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу попасть в личный кабинет", lines=2)
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Как перевести деньги на другую карту?", lines=2)
     btn = gr.Button("Сгенерировать")
 
     gr.Markdown("### ruDialoGPT-small")
@@ -114,4 +109,4 @@ with gr.Blocks() as demo:
         cot3, cot3_time, simple3, simple3_time
     ])
 
-demo.launch()
+demo.launch()