Update app.py

app.py

@@ -3,74 +3,114 @@ import time
 from transformers import pipeline
 from datasets import load_dataset
 
-# Загрузка датасета с отзывами банков
+# Загружаем датасет
 dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
 
-# Отбираем 2 примера для few-shot в CoT
+# Примеры для few-shot (без 'rating')
 few_shot_examples = []
 for row in dataset.select(range(2)):
     review = row["review"]
-    rating = row["rating"]
-    ex = f"Клиент: {review}\nОценка: {rating}\nОтвет: Пожалуйста, разъясните ситуацию в деталях. Мы поможем."
+    ex = f"Клиент: {review}\nОтвет: Спасибо за обращение! Уточните, пожалуйста, детали ситуации, чтобы мы могли помочь."
     few_shot_examples.append(ex)
 
-# Системная инструкция
-system_instruction = (
-    "Ты — вежливый и точный банковский помощник. "
-    "Ты читаешь обращения клиентов и даешь корректные, подробные, официальные ответы. "
-    "Если данных недостаточно — просишь уточнение. Используй рассуждение шаг за шагом."
+# Системные инструкции
+cot_instruction = (
+    "Ты — вежливый банковский помощник. Клиент описывает проблему. "
+    "Сначала проанализируй её шаг за шагом, потом сформулируй итоговый ответ."
 )
 
-# Загружаем три модели
+simple_instruction = (
+    "Ты — вежливый банковский помощник. Отвечай кратко и официально, без лишних деталей. "
+    "Не используй рассуждение, просто дай понятный клиенту ответ."
+)
+
+# Модели
 models = {
     "ruDialoGPT-small": pipeline("text-generation", model="t-bank-ai/ruDialoGPT-small", tokenizer="t-bank-ai/ruDialoGPT-small", device=-1),
     "ruDialoGPT-medium": pipeline("text-generation", model="t-bank-ai/ruDialoGPT-medium", tokenizer="t-bank-ai/ruDialoGPT-medium", device=-1),
     "ruGPT3-small": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_based_on_gpt2", device=-1),
 }
 
-# Построение CoT-промпта
-def build_prompt(user_input):
+# Промпт 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"
+        "Рассуждение и ответ:"
+    )
+
+# Промпт простой
+def build_simple_prompt(user_input):
     examples = "\n\n".join(few_shot_examples)
     return (
-        f"{system_instruction}\n\n"
-        f"{examples}\n\n"
-        f"Клиент: {user_input}\n"
-        f"Опиши шаг за шагом размышления, затем сформулируй окончательный ответ клиенту:"
+        f"{simple_instruction}\n\n{examples}\n\nКлиент: {user_input}\n"
+        "Ответ:"
     )
 
-# Генерация ответов
-def generate_answers(prompt):
+# Генерация ответов по двум промптам
+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():
-        start = time.time()
-        out = pipe(prompt, max_length=300, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
-        elapsed = round(time.time() - start, 2)
-        final_line = out.strip().split('\n')[-1]
-        results[name] = {"answer": final_line, "time": elapsed}
-    return results
-
-# Формат вывода
-def run_models(user_input):
-    prompt = build_prompt(user_input)
-    res = generate_answers(prompt)
+        # CoT
+        start_cot = time.time()
+        out_cot = pipe(prompt_cot, max_length=300, 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]
+
+        # Simple
+        start_simple = time.time()
+        out_simple = pipe(prompt_simple, max_length=300, 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]
+
+        results[name] = {
+            "cot_answer": answer_cot,
+            "cot_time": end_cot,
+            "simple_answer": answer_simple,
+            "simple_time": end_simple
+        }
+
     return (
-        res["ruDialoGPT-small"]["answer"], f"{res['ruDialoGPT-small']['time']} сек",
-        res["ruDialoGPT-medium"]["answer"], f"{res['ruDialoGPT-medium']['time']} сек",
-        res["ruGPT3-small"]["answer"], f"{res['ruGPT3-small']['time']} сек",
+        results["ruDialoGPT-small"]["cot_answer"], f"{results['ruDialoGPT-small']['cot_time']} сек",
+        results["ruDialoGPT-small"]["simple_answer"], f"{results['ruDialoGPT-small']['simple_time']} сек",
+        results["ruDialoGPT-medium"]["cot_answer"], f"{results['ruDialoGPT-medium']['cot_time']} сек",
+        results["ruDialoGPT-medium"]["simple_answer"], f"{results['ruDialoGPT-medium']['simple_time']} сек",
+        results["ruGPT3-small"]["cot_answer"], f"{results['ruGPT3-small']['cot_time']} сек",
+        results["ruGPT3-small"]["simple_answer"], f"{results['ruGPT3-small']['simple_time']} сек",
     )
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 Банковский помощник: CoT + 3 модели (русский язык)")
-    inp = gr.Textbox(label="Запрос клиента", placeholder="Например: Я не могу попасть в личный кабинет", lines=2)
-    btn = gr.Button("Сгенерировать ответы")
-    out1 = gr.Textbox(label="ruDialoGPT-small")
-    t1 = gr.Textbox(label="Время")
-    out2 = gr.Textbox(label="ruDialoGPT-medium")
-    t2 = gr.Textbox(label="Время")
-    out3 = gr.Textbox(label="ruGPT3-small")
-    t3 = gr.Textbox(label="Время")
-
-    btn.click(run_models, inputs=[inp], outputs=[out1, t1, out2, t2, out3, t3])
+    gr.Markdown("## 🏦 Банковский помощник (2 промпта: рассуждение + краткий ответ)")
+
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу попасть в личный кабинет", lines=2)
+    btn = gr.Button("Сгенерировать")
+
+    gr.Markdown("### ruDialoGPT-small")
+    cot1 = gr.Textbox(label="CoT ответ")
+    cot1_time = gr.Textbox(label="Время CoT")
+    simple1 = gr.Textbox(label="Обычный ответ")
+    simple1_time = gr.Textbox(label="Время обычного")
+
+    gr.Markdown("### ruDialoGPT-medium")
+    cot2 = gr.Textbox(label="CoT ответ")
+    cot2_time = gr.Textbox(label="Время CoT")
+    simple2 = gr.Textbox(label="Обычный ответ")
+    simple2_time = gr.Textbox(label="Время обычного")
+
+    gr.Markdown("### ruGPT3-small")
+    cot3 = gr.Textbox(label="CoT ответ")
+    cot3_time = gr.Textbox(label="Время CoT")
+    simple3 = gr.Textbox(label="Обычный ответ")
+    simple3_time = gr.Textbox(label="Время обычного")
+
+    btn.click(generate_dual_answers, inputs=[inp], outputs=[
+        cot1, cot1_time, simple1, simple1_time,
+        cot2, cot2_time, simple2, simple2_time,
+        cot3, cot3_time, simple3, simple3_time
+    ])
 
 demo.launch()