Update app.py

app.py

@@ -3,52 +3,52 @@ import time
 from transformers import pipeline
 from datasets import load_dataset
 
-# Загружаем датасет
+# Загружаем банковский датасет
 dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
 
-# Примеры для few-shot (без 'rating')
+# Примеры для few-shot
 few_shot_examples = []
 for row in dataset.select(range(2)):
     review = row["review"]
-    ex = f"Клиент: {review}\nОтвет: Спасибо за обращение! Уточните, пожалуйста, детали ситуации, чтобы мы могли помочь."
-    few_shot_examples.append(ex)
+    example = f"Клиент: {review}\nКлассификация: прочее"
+    few_shot_examples.append(example)
 
-# Системные инструкции
+# Инструкции
 cot_instruction = (
-    "Ты — вежливый банковский помощник. Клиент описывает проблему."
-    " Проанализируй обращение пошагово: 1. Что случилось? 2. Почему это могло произойти? 3. Как клиенту поступить?"
-    " Вывод: укажи категорию обращения (например: доступ, безопасность, платежи, перевод и т.д.)"
+    "Ты — ассистент банка. Проанализируй обращение клиента и классифицируй его по теме."
+    " Сначала рассуждай шаг за шагом, затем выведи финальную категорию."
 )
 
 simple_instruction = (
-    "Ты — вежливый банковский помощник. Определи кратко, к какой категории относится обращение клиента"
-    " (например: доступ, платежи, безопасность и т.д.)."
+    "Ты — банковский помощник. Классифицируй обращение клиента одним словом — категорией."
 )
 
-# Лёгкие и реальные модели, подходящие под требования
-models = {
-    "ChatGPT-like (ruGPT3-medium)": pipeline("text-generation", model="cointegrated/rugpt3-medium", tokenizer="cointegrated/rugpt3-medium", device=-1),
-    "DeepSeek-like (ruGPT3-small)": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_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):
     examples = "\n\n".join(few_shot_examples)
     return (
         f"{cot_instruction}\n\n{examples}\n\nКлиент: {user_input}\n"
-        "Рассуждение и вывод:"
+        f"Рассуждение:"
     )
 
-# Промпт простой
+
 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"
-        "Категория:"
+        f"Категория:"
     )
 
-# Генерация ответов по двум промптам
+# Рабочие модели с поддержкой русского языка и легкие
+models = {
+    "ChatGPT-like (FRED-T5-small)": pipeline("text2text-generation", model="cointegrated/translation-t5-russian-finetuned", tokenizer="cointegrated/translation-t5-russian-finetuned", device=-1),
+    "DeepSeek-like (ruGPT3-small)": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_based_on_gpt2", device=-1),
+    "GigaChat-like (RuBERT-tiny2)": pipeline("text-classification", model="cointegrated/rubert-tiny2", tokenizer="cointegrated/rubert-tiny2", device=-1),
+}
+
+# Генерация ответов
+
 def generate_dual_answers(user_input):
     results = {}
     prompt_cot = build_cot_prompt(user_input)
@@ -57,62 +57,63 @@ def generate_dual_answers(user_input):
     for name, pipe in models.items():
         # CoT
         start_cot = time.time()
-        out_cot = pipe(prompt_cot, max_new_tokens=200, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        try:
+            out_cot = pipe(prompt_cot, max_new_tokens=150, do_sample=True, top_p=0.9, temperature=0.7)[0]
+            answer_cot = out_cot.get("generated_text", out_cot.get("label", "-"))
+        except:
+            answer_cot = "Ошибка в CoT"
         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_new_tokens=150, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        try:
+            out_simple = pipe(prompt_simple, max_new_tokens=150, do_sample=True, top_p=0.9, temperature=0.7)[0]
+            answer_simple = out_simple.get("generated_text", out_simple.get("label", "-"))
+        except:
+            answer_simple = "Ошибка в обычном"
         end_simple = round(time.time() - start_simple, 2)
-        answer_simple = out_simple.strip().split('\n')[-1]
 
         results[name] = {
-            "cot_answer": answer_cot,
+            "cot_answer": answer_cot.strip(),
             "cot_time": end_cot,
-            "simple_answer": answer_simple,
+            "simple_answer": answer_simple.strip(),
             "simple_time": end_simple
         }
 
     return (
-        results["ChatGPT-like (ruGPT3-medium)"]["cot_answer"], f"{results['ChatGPT-like (ruGPT3-medium)']['cot_time']} сек",
-        results["ChatGPT-like (ruGPT3-medium)"]["simple_answer"], f"{results['ChatGPT-like (ruGPT3-medium)']['simple_time']} сек",
+        results["ChatGPT-like (FRED-T5-small)"]["cot_answer"], f"{results['ChatGPT-like (FRED-T5-small)']['cot_time']} сек",
+        results["ChatGPT-like (FRED-T5-small)"]["simple_answer"], f"{results['ChatGPT-like (FRED-T5-small)']['simple_time']} сек",
+
         results["DeepSeek-like (ruGPT3-small)"]["cot_answer"], f"{results['DeepSeek-like (ruGPT3-small)']['cot_time']} сек",
         results["DeepSeek-like (ruGPT3-small)"]["simple_answer"], f"{results['DeepSeek-like (ruGPT3-small)']['simple_time']} сек",
-        results["GigaChat-like (rubert-tiny2)"]["cot_answer"], f"{results['GigaChat-like (rubert-tiny2)']['cot_time']} сек",
-        results["GigaChat-like (rubert-tiny2)"]["simple_answer"], f"{results['GigaChat-like (rubert-tiny2)']['simple_time']} сек"
+
+        results["GigaChat-like (RuBERT-tiny2)"]["cot_answer"], f"{results['GigaChat-like (RuBERT-tiny2)']['cot_time']} сек",
+        results["GigaChat-like (RuBERT-tiny2)"]["simple_answer"], f"{results['GigaChat-like (RuBERT-tiny2)']['simple_time']} сек",
     )
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏦 Сравнение моделей: Классификация клиентских обращений (CoT и обычный)")
-
-    inp = gr.Textbox(label="Обращение клиента", placeholder="Например: Не могу войти в приложение", lines=2)
-    btn = gr.Button("Получить ответы")
+    gr.Markdown("## 🤖 Классификация клиентских обращений — CoT vs обычный промпт")
+    inp = gr.Textbox(label="Обращение клиента", placeholder="Например: Я не могу войти в личный кабинет", lines=2)
+    btn = gr.Button("Классифицировать")
 
-    gr.Markdown("### ChatGPT-like (ruGPT3-medium)")
-    cot1 = gr.Textbox(label="CoT ответ")
-    cot1_time = gr.Textbox(label="Время CoT")
-    simple1 = gr.Textbox(label="Обычный ответ")
-    simple1_time = gr.Textbox(label="Время обычного")
+    gr.Markdown("### ChatGPT-like (FRED-T5-small)")
+    cot1, cot1_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
+    simple1, simple1_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время обычного")
 
     gr.Markdown("### DeepSeek-like (ruGPT3-small)")
-    cot2 = gr.Textbox(label="CoT ответ")
-    cot2_time = gr.Textbox(label="Время CoT")
-    simple2 = gr.Textbox(label="Обычный ответ")
-    simple2_time = gr.Textbox(label="Время обычного")
+    cot2, cot2_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
+    simple2, simple2_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время обычного")
 
-    gr.Markdown("### GigaChat-like (rubert-tiny2)")
-    cot3 = gr.Textbox(label="CoT ответ")
-    cot3_time = gr.Textbox(label="Время CoT")
-    simple3 = gr.Textbox(label="Обычный ответ")
-    simple3_time = gr.Textbox(label="Время обычного")
+    gr.Markdown("### GigaChat-like (RuBERT-tiny2)")
+    cot3, cot3_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
+    simple3, simple3_time = gr.Textbox(label="Обычный ответ"), 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
+        cot3, cot3_time, simple3, simple3_time,
     ])
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     demo.launch()