Update app.py

app.py

@@ -2,104 +2,105 @@ import gradio as gr
 import time
 from transformers import pipeline
 
-# Проверенные лёгкие русскоязычные модели, работающие в Hugging Face Spaces
+# Подключаем модели (проверенные, лёгкие и с поддержкой русского)
 models = {
-    "ChatGPT-like (ruGPT3small)": pipeline("text-generation", model="ai-forever/rugpt3small_based_on_gpt2", tokenizer="ai-forever/rugpt3small_based_on_gpt2", device=-1),
-    "DeepSeek-like (ruGPT3large)": pipeline("text-generation", model="ai-forever/rugpt3large_based_on_gpt2", tokenizer="ai-forever/rugpt3large_based_on_gpt2", device=-1),
-    "GigaChat-like (rubert-tiny2)": pipeline("text-generation", model="cointegrated/rubert-tiny2", tokenizer="cointegrated/rubert-tiny2", device=-1),
+    "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),
 }
 
 # Промпты
-cot_instruction = """Ты — банковский помощник. Клиент описал проблему.
-Проанализируй обращение шаг за шагом:
-1. Что случилось?
-2. Почему это могло произойти?
-3. Как клиенту поступить?
-Вывод: укажи категорию обращения (например: доступ, безопасность, платежи, перевод и т.д.)"""
-
-simple_instruction = "Ты — банковский помощник. Определи кратко, к какой категории относится обращение клиента (например: доступ, платежи, безопасность и т.д.)."
-
-# Формирование промптов
-
-def build_prompt(instruction, user_input):
-    return f"{instruction}\n\nКлиент: {user_input}"
-
-# Генерация вывода с защитой
-
-def get_output(pipe, prompt, max_tokens=200):
-    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_comparison(user_input):
-    result = {}
-    prompt_cot = build_prompt(cot_instruction, user_input)
-    prompt_simple = build_prompt(simple_instruction, user_input)
+cot_instruction = (
+    "Ты — банковский помощник. Клиент описал проблему. Проанализируй шаг за шагом:\n"
+    "1. Что произошло?\n"
+    "2. Возможные причины?\n"
+    "3. Что делать клиенту?\n"
+    "В конце выдай: Категория обращения (доступ, платежи, безопасность, перевод и т.д.).\n"
+    "\nЗапрос клиента: {user_input}\n"
+)
+
+simple_instruction = (
+    "Ты — банковский помощник. Определи, к какой категории относится обращение клиента: доступ, платежи, безопасность, перевод и т.д.\n"
+    "Запрос клиента: {user_input}\n"
+    "Категория:"
+)
+
+# Функции построения промптов
+def build_cot_prompt(user_input):
+    return cot_instruction.format(user_input=user_input)
+
+def build_simple_prompt(user_input):
+    return simple_instruction.format(user_input=user_input)
+
+# Генерация
+
+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 ответ
+        # CoT
         start_cot = time.time()
-        answer_cot = get_output(pipe, prompt_cot)
-        time_cot = round(time.time() - start_cot, 2)
+        out_cot = pipe(prompt_cot, max_length=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()
-        answer_simple = get_output(pipe, prompt_simple)
-        time_simple = round(time.time() - start_simple, 2)
-
-        result[name] = {
-            "cot_answer": answer_cot.strip(),
-            "cot_time": f"{time_cot} сек",
-            "simple_answer": answer_simple.strip(),
-            "simple_time": f"{time_simple} сек"
+        out_simple = pipe(prompt_simple, max_length=200, 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 (
-        result["ChatGPT-like (ruGPT3small)"]["cot_answer"], result["ChatGPT-like (ruGPT3small)"]["cot_time"],
-        result["ChatGPT-like (ruGPT3small)"]["simple_answer"], result["ChatGPT-like (ruGPT3small)"]["simple_time"],
-        result["DeepSeek-like (ruGPT3large)"]["cot_answer"], result["DeepSeek-like (ruGPT3large)"]["cot_time"],
-        result["DeepSeek-like (ruGPT3large)"]["simple_answer"], result["DeepSeek-like (ruGPT3large)"]["simple_time"],
-        result["GigaChat-like (rubert-tiny2)"]["cot_answer"], result["GigaChat-like (rubert-tiny2)"]["cot_time"],
-        result["GigaChat-like (rubert-tiny2)"]["simple_answer"], result["GigaChat-like (rubert-tiny2)"]["simple_time"]
+        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']} сек",
     )
 
 # Интерфейс
 with gr.Blocks() as demo:
-    gr.Markdown("## Сравнение моделей: ruGPT3small, ruGPT3large, rubert-tiny2 (Классификация обращений)")
+    gr.Markdown("## Сравнение моделей: ruGPT3small, saiga2, ruBERT (Классификация обращений)")
 
-    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу войти в личный кабинет", lines=2)
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу оплатить картой в магазине, пишет «техническая ошибка»")
     btn = gr.Button("Сгенерировать")
 
-    # ChatGPT-like
-    gr.Markdown("### ChatGPT-like (ruGPT3small)")
+    gr.Markdown("### ChatGPT-like (saiga2_7b_lora)")
     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 (ruGPT3large)")
+    gr.Markdown("### DeepSeek-like (ruGPT3small)")
     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 (rubert-tiny2)")
+    gr.Markdown("### GigaChat-like (rubert-base-sentiment)")
     cot3 = gr.Textbox(label="CoT ответ")
     cot3_time = gr.Textbox(label="Время CoT")
     simple3 = gr.Textbox(label="Обычный ответ")
     simple3_time = gr.Textbox(label="Время обычного")
 
-    btn.click(generate_comparison, inputs=[inp], outputs=[
+    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
     ])
 
+
 if __name__ == '__main__':
     demo.launch()