Update app.py

app.py

@@ -6,41 +6,44 @@ from datasets import load_dataset
 # Загружаем датасет
 dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
 
-# Примеры классификации (вручную или через разметку датасета)
-few_shot_examples = [
-    "Клиент: Я не могу войти в личный кабинет\nКлассификация: Проблема с доступом",
-    "Клиент: Хочу оформить кредит на авто\nКлассификация: Кредитование",
-    "Клиент: Почему списали деньги с карты?\nКлассификация: Жалоба на транзакцию"
-]
+# Примеры для few-shot
+few_shot_examples = []
+for row in dataset.select(range(3)):
+    review = row["review"]
+    category = row["category"] if "category" in row else "(Категория)"
+    ex = f"Клиент: {review}\nКлассификация: {category}"
+    few_shot_examples.append(ex)
 
 # Инструкции
 cot_instruction = (
-    "Ты — банковский помощник. Клиент описывает обращение."
-    " Проанализируй обращение пошагово, определи его суть и укажи категорию обращения."
-    " Дай только итоговую классификацию."
+    "Ты — помощник банка. Клиент задал вопрос. Проанализируй обращение шаг за шагом, "
+    "выдели ключевые признаки и выдай итоговую категорию обращения."
 )
 
 simple_instruction = (
-    "Ты — банковский помощник. Определи, к какой категории относится обращение клиента."
-    " Ответ должен быть кратким: только категория."
+    "Ты — помощник банка. Определи категорию обращения клиента. Ответ должен быть кратким, без лишнего текста."
 )
 
-# Промпт CoT
+# Используемые модели
+models = {
+    "ChatGPT-like (ruGPT3small)": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_based_on_gpt2", device=-1),
+    "GigaChat-like (ruDialoGPT-medium)": pipeline("text-generation", model="tinkoff-ai/ruDialoGPT-medium", tokenizer="tinkoff-ai/ruDialoGPT-medium", device=-1),
+    "DeepSeek-like (RuBERT-tiny2)": pipeline("text-classification", model="cointegrated/rubert-tiny2", tokenizer="cointegrated/rubert-tiny2", device=-1)
+}
+
+# Формирование промптов
+
 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Классификация:"
-
-# Подключаем реальные модели (только существующие и совместимые)
-models = {
-    "GPT2-large": pipeline("text-generation", model="cointegrated/rugpt2-large", tokenizer="cointegrated/rugpt2-large", device=-1),
-    "RuBERT-tiny2": pipeline("text-classification", model="cointegrated/rubert-tiny2", tokenizer="cointegrated/rubert-tiny2", device=-1),
-    "ruGPT3-medium": pipeline("text-generation", model="IlyaGusev/rugpt3medium_based_on_gpt2", tokenizer="IlyaGusev/rugpt3medium_based_on_gpt2", device=-1),
-}
+    return (
+        f"{simple_instruction}\n\n{examples}\n\nКлиент: {user_input}\nКлассификация:"
+    )
 
 # Генерация ответов
 
@@ -50,57 +53,70 @@ def generate_dual_answers(user_input):
     prompt_simple = build_simple_prompt(user_input)
 
     for name, pipe in models.items():
-        if "classification" in str(pipe.task):
+        if name.startswith("DeepSeek"):
+            # классификация
             start = time.time()
-            simple = pipe(user_input)[0]['label']
+            output = pipe(user_input)[0]
             end = round(time.time() - start, 2)
             results[name] = {
-                "cot": "(не поддерживается)",
-                "cot_time": "-",
-                "simple": simple,
-                "simple_time": f"{end} сек"
+                "cot_answer": output['label'],
+                "cot_time": end,
+                "simple_answer": output['label'],
+                "simple_time": end
             }
         else:
+            # генерация CoT
             start_cot = time.time()
-            out_cot = pipe(prompt_cot, max_length=200, 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.9, temperature=0.7)[0]["generated_text"]
             end_cot = round(time.time() - start_cot, 2)
-            cot_answer = out_cot.split("Классификация:")[-1].strip()
+            answer_cot = out_cot.split("Классификация:")[-1].strip()
 
+            # генерация 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=60, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
             end_simple = round(time.time() - start_simple, 2)
-            simple_answer = out_simple.split("Классификация:")[-1].strip()
+            answer_simple = out_simple.split("Классификация:")[-1].strip()
 
             results[name] = {
-                "cot": cot_answer,
-                "cot_time": f"{end_cot} сек",
-                "simple": simple_answer,
-                "simple_time": f"{end_simple} сек"
+                "cot_answer": answer_cot,
+                "cot_time": end_cot,
+                "simple_answer": answer_simple,
+                "simple_time": end_simple
             }
 
     return (
-        results["GPT2-large"]["cot"], results["GPT2-large"]["cot_time"], results["GPT2-large"]["simple"], results["GPT2-large"]["simple_time"],
-        results["RuBERT-tiny2"]["cot"], results["RuBERT-tiny2"]["cot_time"], results["RuBERT-tiny2"]["simple"], results["RuBERT-tiny2"]["simple_time"],
-        results["ruGPT3-medium"]["cot"], results["ruGPT3-medium"]["cot_time"], results["ruGPT3-medium"]["simple"], results["ruGPT3-medium"]["simple_time"]
+        results["ChatGPT-like (ruGPT3small)"]["cot_answer"], f"{results['ChatGPT-like (ruGPT3small)']['cot_time']} сек",
+        results["ChatGPT-like (ruGPT3small)"]["simple_answer"], f"{results['ChatGPT-like (ruGPT3small)']['simple_time']} сек",
+        results["GigaChat-like (ruDialoGPT-medium)"]["cot_answer"], f"{results['GigaChat-like (ruDialoGPT-medium)']['cot_time']} сек",
+        results["GigaChat-like (ruDialoGPT-medium)"]["simple_answer"], f"{results['GigaChat-like (ruDialoGPT-medium)']['simple_time']} сек",
+        results["DeepSeek-like (RuBERT-tiny2)"]["cot_answer"], f"{results['DeepSeek-like (RuBERT-tiny2)']['cot_time']} сек",
+        results["DeepSeek-like (RuBERT-tiny2)"]["simple_answer"], f"{results['DeepSeek-like (RuBERT-tiny2)']['simple_time']} сек"
     )
 
-# Интерфейс Gradio
+# Gradio интерфейс
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏦 Классификация клиентских обращений (CoT и обычный промпт)")
-    inp = gr.Textbox(label="Обращение клиента", placeholder="Например: Почему не работает мобильный банк?", lines=2)
-    btn = gr.Button("Анализировать")
-
-    gr.Markdown("### GPT2-large")
-    cot1, cot1_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
-    simple1, simple1_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время ответа")
-
-    gr.Markdown("### RuBERT-tiny2")
-    cot2, cot2_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
-    simple2, simple2_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время ответа")
-
-    gr.Markdown("### ruGPT3-medium")
-    cot3, cot3_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
-    simple3, simple3_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время ответа")
+    gr.Markdown("## 🏦 Классификация клиентских обращений — Сравнение моделей и промптов")
+
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Не приходит СМС-код для входа в приложение", lines=2)
+    btn = gr.Button("Классифицировать")
+
+    gr.Markdown("### ChatGPT-like (ruGPT3small)")
+    cot1 = gr.Textbox(label="CoT ответ")
+    cot1_time = gr.Textbox(label="Время CoT")
+    simple1 = gr.Textbox(label="Обычный ответ")
+    simple1_time = gr.Textbox(label="Время обычного")
+
+    gr.Markdown("### GigaChat-like (ruDialoGPT-medium)")
+    cot2 = gr.Textbox(label="CoT ответ")
+    cot2_time = gr.Textbox(label="Время CoT")
+    simple2 = gr.Textbox(label="Обычный ответ")
+    simple2_time = gr.Textbox(label="Время обычного")
+
+    gr.Markdown("### DeepSeek-like (RuBERT-tiny2)")
+    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,