Update app.py

app.py

@@ -1,39 +1,54 @@
 import gradio as gr
 import time
 from transformers import pipeline
+from datasets import load_dataset
 
-# Подключаем модели (проверенные, лёгкие и с поддержкой русского)
-models = {
-    "ChatGPT-like (saiga2_7b_lora)": pipeline("text2text-generation", model="IlyaGusev/saiga2_7b_lora", tokenizer="IlyaGusev/saiga2_7b_lora", device=-1),
-    "DeepSeek-like (ruGPT3small)": pipeline("text-generation", model="ai-forever/ruGPT3Small", tokenizer="ai-forever/ruGPT3Small", device=-1),
-    "GigaChat-like (rubert-base-sentiment)": pipeline("text2text-generation", model="blanchefort/rubert-base-cased-sentiment", tokenizer="blanchefort/rubert-base-cased-sentiment", device=-1),
-}
+# Загружаем датасет
+dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
 
-# Промпты
+# Примеры для 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 = (
-    "Ты — банковский помощник. Клиент описал проблему. Проанализируй шаг за шагом:\n"
-    "1. Что произошло?\n"
-    "2. Возможные причины?\n"
-    "3. Что делать клиенту?\n"
-    "В конце выдай: Категория обращения (доступ, платежи, безопасность, перевод и т.д.).\n"
-    "\nЗапрос клиента: {user_input}\n"
+    "Ты — вежливый банковский помощник. Клиент описывает проблему."
+    " Проанализируй обращение пошагово: 1. Что случилось? 2. Почему это могло произойти? 3. Как клиенту поступить?"
+    " Вывод: укажи категорию обращения (например: доступ, безопасность, платежи, перевод и т.д.)"
 )
 
 simple_instruction = (
-    "Ты — банковский помощник. Определи, к какой категории относится обращение клиента: доступ, платежи, безопасность, перевод и т.д.\n"
-    "Запрос клиента: {user_input}\n"
-    "Категория:"
+    "Ты — вежливый банковский помощник. Определи кратко, к какой категории относится обращение клиента"
+    " (например: доступ, платежи, безопасность и т.д.)."
 )
 
-# Функции построения промптов
+# Модели
+models = {
+    "ChatGPT-like (saiga_mistral)": pipeline("text-generation", model="IlyaGusev/saiga_mistral_7b_merged", tokenizer="IlyaGusev/saiga_mistral_7b_merged", device=-1),
+    "DeepSeek-like (ruGPT3-medium)": pipeline("text-generation", model="ai-forever/rugpt3medium_based_on_gpt2", tokenizer="ai-forever/rugpt3medium_based_on_gpt2", device=-1),
+    "GigaChat-like (rubert-tiny2)": pipeline("text-generation", model="cointegrated/rubert-tiny2", tokenizer="cointegrated/rubert-tiny2", device=-1),
+}
+
+# Промпт CoT
 def build_cot_prompt(user_input):
-    return cot_instruction.format(user_input=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):
-    return simple_instruction.format(user_input=user_input)
-
-# Генерация
+    examples = "\n\n".join(few_shot_examples)
+    return (
+        f"{simple_instruction}\n\n{examples}\n\nКлиент: {user_input}\n"
+        "Категория:"
+    )
 
+# Генерация ответов по двум промптам
 def generate_dual_answers(user_input):
     results = {}
     prompt_cot = build_cot_prompt(user_input)
@@ -42,13 +57,13 @@ def generate_dual_answers(user_input):
     for name, pipe in models.items():
         # 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"]
+        out_cot = pipe(prompt_cot, max_new_tokens=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=200, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        out_simple = pipe(prompt_simple, max_new_tokens=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]
 
@@ -59,37 +74,40 @@ def generate_dual_answers(user_input):
             "simple_time": end_simple
         }
 
-    return (
-        results["ChatGPT-like (saiga2_7b_lora)"]["cot_answer"], f"{results['ChatGPT-like (saiga2_7b_lora)']['cot_time']} сек",
-        results["ChatGPT-like (saiga2_7b_lora)"]["simple_answer"], f"{results['ChatGPT-like (saiga2_7b_lora)']['simple_time']} сек",
-
-        results["DeepSeek-like (ruGPT3small)"]["cot_answer"], f"{results['DeepSeek-like (ruGPT3small)']['cot_time']} сек",
-        results["DeepSeek-like (ruGPT3small)"]["simple_answer"], f"{results['DeepSeek-like (ruGPT3small)']['simple_time']} сек",
-
-        results["GigaChat-like (rubert-base-sentiment)"]["cot_answer"], f"{results['GigaChat-like (rubert-base-sentiment)']['cot_time']} сек",
-        results["GigaChat-like (rubert-base-sentiment)"]["simple_answer"], f"{results['GigaChat-like (rubert-base-sentiment)']['simple_time']} сек",
+    return tuple(
+        [
+            results[m]["cot_answer"], f"{results[m]['cot_time']} сек",
+            results[m]["simple_answer"], f"{results[m]['simple_time']} сек"
+        ]
+        for m in models
+    )[0] + tuple(
+        [
+            results[m]["cot_answer"], f"{results[m]['cot_time']} сек",
+            results[m]["simple_answer"], f"{results[m]['simple_time']} сек"
+        ]
+        for m in list(models)[1:]
     )
 
-# Интерфейс
+# Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## Сравнение моделей: ruGPT3small, saiga2, ruBERT (Классификация обращений)")
+    gr.Markdown("## 🏦 Сравнение моделей (Классификация клиентских обращений): ChatGPT-like, DeepSeek-like, GigaChat-like")
 
-    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу оплатить картой в магазине, пишет «техническая ошибка»")
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу попасть в личный кабинет", lines=2)
     btn = gr.Button("Сгенерировать")
 
-    gr.Markdown("### ChatGPT-like (saiga2_7b_lora)")
+    gr.Markdown("### ChatGPT-like (saiga_mistral)")
     cot1 = gr.Textbox(label="CoT ответ")
     cot1_time = gr.Textbox(label="Время CoT")
     simple1 = gr.Textbox(label="Обычный ответ")
     simple1_time = gr.Textbox(label="Время обычного")
 
-    gr.Markdown("### DeepSeek-like (ruGPT3small)")
+    gr.Markdown("### DeepSeek-like (ruGPT3-medium)")
     cot2 = gr.Textbox(label="CoT ответ")
     cot2_time = gr.Textbox(label="Время CoT")
     simple2 = gr.Textbox(label="Обычный ответ")
     simple2_time = gr.Textbox(label="Время обычного")
 
-    gr.Markdown("### GigaChat-like (rubert-base-sentiment)")
+    gr.Markdown("### GigaChat-like (rubert-tiny2)")
     cot3 = gr.Textbox(label="CoT ответ")
     cot3_time = gr.Textbox(label="Время CoT")
     simple3 = gr.Textbox(label="Обычный ответ")
@@ -101,6 +119,5 @@ with gr.Blocks() as demo:
         cot3, cot3_time, simple3, simple3_time
     ])
 
-
-if __name__ == '__main__':
+if __name__ == "__main__":
     demo.launch()