Update app.py

app.py

@@ -3,116 +3,127 @@ import time
 from transformers import pipeline
 from datasets import load_dataset
 
-# Загружаем банковский датасет
+# Загружаем датасет
 dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
 
-# Примеры для few-shot
+# Примеры для few-shot (без 'rating')
 few_shot_examples = []
 for row in dataset.select(range(2)):
     review = row["review"]
-    example = f"Клиент: {review}\nКлассификация: прочее"
-    few_shot_examples.append(example)
+    ex = f"Клиент: {review}\nОтвет: Спасибо за обращение! Уточните, пожалуйста, детали ситуации, чтобы мы могли помочь."
+    few_shot_examples.append(ex)
 
-# Инструкции
+# Системные инструкции
 cot_instruction = (
-    "Ты — ассистент банка. Проанализируй обращение клиента и классифицируй его по теме."
-    " Сначала рассуждай шаг за шагом, затем выведи финальную категорию."
+    "Ты — банковский помощник. Твоя задача — классифицировать клиентское обращение.\n"
+    "Проанализируй обращение пошагово, выдели ключевые слова, выясни намерение клиента,\n"
+    "и отнеси его к одной из категорий: вход в ЛК, SMS, заявка, ошибка, перевод, карта, другое."
 )
 
 simple_instruction = (
-    "Ты — банковский помощник. Классифицируй обращение клиента одним словом — категорией."
+    "Ты — банковский помощник. Классифицируй обращение пользователя кратко и по существу,\n"
+    "укажи одну категорию: вход в ЛК, SMS, заявка, ошибка, перевод, карта, другое."
 )
 
-# Промпты
+# Модели
+models = {
+    "ChatGPT-like (ruGPT3-small)": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_based_on_gpt2", device=-1),
+    "DeepSeek-like (rubert-tiny2)": pipeline("text-classification", model="cointegrated/rubert-tiny2", tokenizer="cointegrated/rubert-tiny2", device=-1),
+    "GigaChat-like (sberbank-ai/rugpt3medium_based_on_gpt2)": pipeline("text-generation", model="sberbank-ai/rugpt3medium_based_on_gpt2", tokenizer="sberbank-ai/rugpt3medium_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"
-        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)
     prompt_simple = build_simple_prompt(user_input)
 
     for name, pipe in models.items():
-        # CoT
-        start_cot = time.time()
-        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)
-
-        # Simple
-        start_simple = time.time()
-        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)
-
-        results[name] = {
-            "cot_answer": answer_cot.strip(),
-            "cot_time": end_cot,
-            "simple_answer": answer_simple.strip(),
-            "simple_time": end_simple
-        }
+        if name.startswith("DeepSeek-like"):
+            # Text-classification модель
+            start_simple = time.time()
+            classification = pipe(user_input)[0]['label']
+            end_simple = round(time.time() - start_simple, 2)
+            results[name] = {
+                "cot_answer": "(CoT не поддерживается)",
+                "cot_time": "-",
+                "simple_answer": classification,
+                "simple_time": end_simple
+            }
+        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"]
+            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=150, 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 (
-        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["ChatGPT-like (ruGPT3-small)"]["cot_answer"], f"{results['ChatGPT-like (ruGPT3-small)']['cot_time']} сек",
+        results["ChatGPT-like (ruGPT3-small)"]["simple_answer"], f"{results['ChatGPT-like (ruGPT3-small)']['simple_time']} сек",
+        results["DeepSeek-like (rubert-tiny2)"]["cot_answer"], results["DeepSeek-like (rubert-tiny2)"]["cot_time"],
+        results["DeepSeek-like (rubert-tiny2)"]["simple_answer"], f"{results['DeepSeek-like (rubert-tiny2)']['simple_time']} сек",
+        results["GigaChat-like (sberbank-ai/rugpt3medium_based_on_gpt2)"]["cot_answer"], f"{results['GigaChat-like (sberbank-ai/rugpt3medium_based_on_gpt2)']['cot_time']} сек",
+        results["GigaChat-like (sberbank-ai/rugpt3medium_based_on_gpt2)"]["simple_answer"], f"{results['GigaChat-like (sberbank-ai/rugpt3medium_based_on_gpt2)']['simple_time']} сек",
     )
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 Классификация клиентских обращений — CoT vs обычный промпт")
-    inp = gr.Textbox(label="Обращение клиента", placeholder="Например: Я не могу войти в личный кабинет", lines=2)
+    gr.Markdown("## 🏦 Классификация клиентских обращений (CoT + обычный)")
+
+    inp = gr.Textbox(label="Обращение клиента", placeholder="Например: Я не могу попасть в личный кабинет", lines=2)
     btn = gr.Button("Классифицировать")
 
-    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("### ChatGPT-like (ruGPT3-small)")
+    cot1 = gr.Textbox(label="CoT ответ")
+    cot1_time = gr.Textbox(label="Время CoT")
+    simple1 = gr.Textbox(label="Обычный ответ")
+    simple1_time = gr.Textbox(label="Время обычного")
 
-    gr.Markdown("### DeepSeek-like (ruGPT3-small)")
-    cot2, cot2_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
-    simple2, simple2_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время обычного")
+    gr.Markdown("### DeepSeek-like (rubert-tiny2)")
+    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-tiny2)")
-    cot3, cot3_time = gr.Textbox(label="CoT ответ"), gr.Textbox(label="Время CoT")
-    simple3, simple3_time = gr.Textbox(label="Обычный ответ"), gr.Textbox(label="Время обычного")
+    gr.Markdown("### GigaChat-like (ruGPT3-medium)")
+    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,
+        cot3, cot3_time, simple3, simple3_time
     ])
 
 if __name__ == '__main__':