Update app.py

app.py

@@ -2,108 +2,106 @@ import gradio as gr
 import time
 from transformers import pipeline
 
-# Инициализация моделей
+# Настройка моделей
 models = {
-    "ChatGPT-like": pipeline(
-        "text-generation",
-        model="tiiuae/falcon-rw-1b",
-        tokenizer="tiiuae/falcon-rw-1b",
-        device=-1
-    ),
-    "DeepSeek-like": pipeline(
-        "text-generation",
-        model="deepseek-ai/DeepSeek-Coder-1.3B-instruct",
-        tokenizer="deepseek-ai/DeepSeek-Coder-1.3B-instruct",
-        device=-1
-    ),
-    "GigaChat-like": pipeline(
-        "text-generation",
-        model="tinkoff-ai/ruDialoGPT-medium",
-        tokenizer="tinkoff-ai/ruDialoGPT-medium",
-        device=-1
-    )
+    "ChatGPT-like (FRED-T5)": pipeline("text2text-generation", model="ai-forever/FRED-T5-1.7B", tokenizer="ai-forever/FRED-T5-1.7B", device_map="auto"),
+    "DeepSeek-like (Qwen7B)": pipeline("text-generation", model="lightblue/DeepSeek-R1-Distill-Qwen-7B-Multilingual", tokenizer="lightblue/DeepSeek-R1-Distill-Qwen-7B-Multilingual", device_map="auto"),
+    "GigaChat-like (GigaChat-20B)": pipeline("text-generation", model="ai-sage/GigaChat-20B-A3B-instruct", tokenizer="ai-sage/GigaChat-20B-A3B-instruct", device_map="auto")
 }
 
 # Промпты
-def build_simple_prompt(user_input):
-    return f"Клиент: {user_input}\nКатегория обращения:"
+cot_instruction = (
+    "Ты — банковский специалист. Клиент задал вопрос."
+    " Проанализируй обращение пошагово:
+    1. В чём проблема?
+    2. Возможные причины?
+    3. Как решить?
+    Вывод: категория обращения."
+)
 
-def build_cot_prompt(user_input):
-    return (
-        f"Клиент: {user_input}\n"
-        "Проанализируй обращение пошагово:\n"
-        "1. В чём проблема?\n"
-        "2. Возможные причины?\n"
-        "3. Как решить?\n"
-        "Вывод: категория обращения:"
-    )
+simple_instruction = "Ты — банковский специалист. Определи категорию обращения клиента кратко."
+
+# Генерация промптов
+
+def build_prompt(instruction, user_input):
+    return f"{instruction}\n\nКлиент: {user_input}"
+
+# Обработка ответа
+
+def get_output(pipe, prompt, max_tokens=300):
+    try:
+        output = pipe(prompt, max_new_tokens=max_tokens, truncation=True, do_sample=True, temperature=0.7)[0]
+        return output.get("generated_text") or output.get("output_text") or "(нет ответа)"
+    except Exception as e:
+        return f"Ошибка: {e}"
 
-# Генерация
-def generate_classification(user_input):
-    prompt_simple = build_simple_prompt(user_input)
-    prompt_cot = build_cot_prompt(user_input)
+# Основная функция
 
-    results = {}
+def generate_comparison(user_input):
+    result = {}
+    prompt_cot = build_prompt(cot_instruction, user_input)
+    prompt_simple = build_prompt(simple_instruction, user_input)
 
     for name, pipe in models.items():
-        # CoT
+        # CoT ответ
         start_cot = time.time()
-        cot_out = pipe(prompt_cot, max_new_tokens=120, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        answer_cot = get_output(pipe, prompt_cot)
         time_cot = round(time.time() - start_cot, 2)
 
-        # Simple
+        # Обычный ответ
         start_simple = time.time()
-        simple_out = pipe(prompt_simple, max_new_tokens=60, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
+        answer_simple = get_output(pipe, prompt_simple)
         time_simple = round(time.time() - start_simple, 2)
 
-        results[name] = {
-            "cot": cot_out.strip(),
-            "cot_time": time_cot,
-            "simple": simple_out.strip(),
-            "simple_time": time_simple
+        result[name] = {
+            "cot_answer": answer_cot.strip(),
+            "cot_time": f"{time_cot} сек",
+            "simple_answer": answer_simple.strip(),
+            "simple_time": f"{time_simple} сек"
         }
 
     return (
-        results["ChatGPT-like"]["cot"], f"{results['ChatGPT-like']['cot_time']} сек",
-        results["ChatGPT-like"]["simple"], f"{results['ChatGPT-like']['simple_time']} сек",
-        results["DeepSeek-like"]["cot"], f"{results['DeepSeek-like']['cot_time']} сек",
-        results["DeepSeek-like"]["simple"], f"{results['DeepSeek-like']['simple_time']} сек",
-        results["GigaChat-like"]["cot"], f"{results['GigaChat-like']['cot_time']} сек",
-        results["GigaChat-like"]["simple"], f"{results['GigaChat-like']['simple_time']} сек"
+        result["ChatGPT-like (FRED-T5)"]["cot_answer"], result["ChatGPT-like (FRED-T5)"]["cot_time"],
+        result["ChatGPT-like (FRED-T5)"]["simple_answer"], result["ChatGPT-like (FRED-T5)"]["simple_time"],
+        result["DeepSeek-like (Qwen7B)"]["cot_answer"], result["DeepSeek-like (Qwen7B)"]["cot_time"],
+        result["DeepSeek-like (Qwen7B)"]["simple_answer"], result["DeepSeek-like (Qwen7B)"]["simple_time"],
+        result["GigaChat-like (GigaChat-20B)"]["cot_answer"], result["GigaChat-like (GigaChat-20B)"]["cot_time"],
+        result["GigaChat-like (GigaChat-20B)"]["simple_answer"], result["GigaChat-like (GigaChat-20B)"]["simple_time"]
     )
 
 # Интерфейс
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏦 Сравнение моделей: ChatGPT, DeepSeek, GigaChat (банковская классификация)")
+    gr.Markdown("## Сравнение моделей: ChatGPT, DeepSeek, GigaChat (банковская классификация)")
 
-    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу войти в личный кабинет", lines=2)
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу оплатить картой", lines=2)
     btn = gr.Button("Сгенерировать")
 
     # ChatGPT-like
-    gr.Markdown("### ChatGPT-like (Falcon RW 1B)")
+    gr.Markdown("### ChatGPT-like (FRED-T5)")
     cot1 = gr.Textbox(label="CoT ответ")
     cot1_time = gr.Textbox(label="Время CoT")
     simple1 = gr.Textbox(label="Обычный ответ")
     simple1_time = gr.Textbox(label="Время обычного")
 
     # DeepSeek-like
-    gr.Markdown("### DeepSeek-like")
+    gr.Markdown("### DeepSeek-like (Qwen7B)")
     cot2 = gr.Textbox(label="CoT ответ")
     cot2_time = gr.Textbox(label="Время CoT")
     simple2 = gr.Textbox(label="Обычный ответ")
     simple2_time = gr.Textbox(label="Время обычного")
 
     # GigaChat-like
-    gr.Markdown("### GigaChat-like (ruDialoGPT-medium)")
+    gr.Markdown("### GigaChat-like (GigaChat-20B)")
     cot3 = gr.Textbox(label="CoT ответ")
     cot3_time = gr.Textbox(label="Время CoT")
     simple3 = gr.Textbox(label="Обычный ответ")
     simple3_time = gr.Textbox(label="Время обычного")
 
-    btn.click(generate_classification, inputs=[inp], outputs=[
+    btn.click(generate_comparison, inputs=[inp], outputs=[
         cot1, cot1_time, simple1, simple1_time,
         cot2, cot2_time, simple2, simple2_time,
         cot3, cot3_time, simple3, simple3_time
     ])
 
-demo.launch()
+if __name__ == '__main__':
+    demo.launch()