Update app.py

app.py

@@ -3,74 +3,68 @@ import torch
 import time
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# 1) Конфигурация доступных моделей
+# 1) Публичные русскоязычные модели
 MODEL_CONFIGS = {
-    "GigaChat-like": "ai-forever/rugpt3large_based_on_gpt2",
-    "ChatGPT-like":  "ai-forever/rugpt3medium_based_on_gpt2",
-    "DeepSeek-like": "ai-forever/rugpt3small_based_on_gpt2"
+    "GigaChat-like":  "ai-forever/rugpt3large_based_on_gpt2",
+    "ChatGPT-like":   "ai-forever/rugpt3medium_based_on_gpt2",
+    "DeepSeek-like":  "ai-forever/rugpt3small_based_on_gpt2"
 }
 
-# 2) Выбор устройства
+# 2) Устройство
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 3) Загрузка токенизаторов и моделей один раз при старте
+# 3) Загружаем модели один раз при старте
 models = {}
-for label, repo_id in MODEL_CONFIGS.items():
-    try:
-        tok = AutoTokenizer.from_pretrained(repo_id)
-        mdl = AutoModelForCausalLM.from_pretrained(repo_id)
-        mdl.to(device).eval()
-        models[label] = (tok, mdl)
-    except Exception as e:
-        print(f"Не удалось загрузить {repo_id}: {e}")
+for name, repo_id in MODEL_CONFIGS.items():
+    tok = AutoTokenizer.from_pretrained(repo_id)
+    mdl = AutoModelForCausalLM.from_pretrained(repo_id)
+    mdl.to(device).eval()
+    models[name] = (tok, mdl)
 
-# 4) Chain-of-Thought промпты
-def cot_prompt_1(q): return f"Клиент: «{q}»\nШаг за шагом объясни ответ от лица банка."
-def cot_prompt_2(q): return f"Клиент: «{q}»\nРазбери вопрос на части и дай развёрнутый ответ."
+# 4) CoT-промпты
+def cot1(q): return f"Клиент: «{q}»\nШаг за шагом объясни, как ответил бы банк."
+def cot2(q): return f"Клиент: «{q}»\nРазбери запрос и дай развернутый ответ."
 
-# 5) Функция генерации
-def generate_all_responses(question):
-    if not question.strip():
-        return {k: {"error": "Пустой вопрос"} for k in models}
+# 5) Генерация ответов + замер времени
+def generate_all(q):
     out = {}
     for name, (tok, mdl) in models.items():
         out[name] = {}
-        for idx, prm in enumerate((cot_prompt_1, cot_prompt_2), start=1):
-            prompt = prm(question)
-            try:
-                inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
-                t0 = time.time()
-                with torch.no_grad():
-                    ids = mdl.generate(**inputs, max_new_tokens=150, do_sample=True, temperature=0.7, top_p=0.9)
-                t1 = time.time()
-                txt = tok.decode(ids[0], skip_special_tokens=True)
-                if txt.startswith(prompt): txt = txt[len(prompt):].strip()
-                out[name][f"CoT-промпт {idx}"] = {
-                    "response": txt or "— пустой ответ —",
-                    "time": f"{round(t1-t0,2)} сек."
-                }
-            except Exception as e:
-                out[name][f"CoT-промпт {idx}"] = {"response": f"Ошибка генерации: {e}", "time": "-"}
+        for idx, prm in enumerate((cot1, cot2), start=1):
+            prompt = prm(q)
+            inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
+            t0 = time.time()
+            with torch.no_grad():
+                ids = mdl.generate(**inputs, max_new_tokens=150, do_sample=True, temperature=0.7, top_p=0.9)
+            dt = round(time.time() - t0, 2)
+            resp = tok.decode(ids[0], skip_special_tokens=True)
+            if resp.startswith(prompt):
+                resp = resp[len(prompt):].strip()
+            out[name][f"CoT-промпт {idx}"] = f"{resp}\n⏱ {dt} сек."
     return out
 
 # 6) Обёртка для Gradio
-def run_all(question):
-    res = generate_all_responses(question)
-    md = []
-    for model_name, prompts in res.items():
-        md.append(f"### 🔹 {model_name}")
-        for label, data in prompts.items():
-            md.append(f"**{label}** ({data['time']}):\n> {data['response']}")
-    return "\n\n".join(md)
+def run_all(q):
+    res = generate_all(q)
+    # вернём 3 больших текста: сначала GigaChat-like, потом ChatGPT-like, потом DeepSeek-like
+    return (
+        "\n\n".join(f"### {k}\n\n" + "\n\n".join(v.values()) for k, v in [("GigaChat-like", res["GigaChat-like"])]),
+        "\n\n".join(f"### {k}\n\n" + "\n\n".join(v.values()) for k, v in [("ChatGPT-like", res["ChatGPT-like"])]),
+        "\n\n".join(f"### {k}\n\n" + "\n\n".join(v.values()) for k, v in [("DeepSeek-like", res["DeepSeek-like"])]),
+    )
 
-# 7) Интерфейс Gradio с блоками
+# 7) Blocks-интерфейс с явным полем вывода
 with gr.Blocks() as demo:
     gr.Markdown("# Alpha Bank Assistant — сравнение CoT-моделей")
-    inp = gr.Textbox(lines=3, placeholder="Введите вопрос клиента...", label="Вопрос клиента")
+    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Как восстановить доступ в мобильный банк?", lines=3)
     btn = gr.Button("Сгенерировать ответы")
-    out = gr.Markdown(label="Результаты")
-    btn.click(fn=run_all, inputs=inp, outputs=out)
+    # вот поле вывода: три текстовых Textbox’а под кнопкой
+    out1 = gr.Textbox(label="GigaChat-like", lines=8)
+    out2 = gr.Textbox(label="ChatGPT-like",  lines=8)
+    out3 = gr.Textbox(label="DeepSeek-like", lines=8)
+    btn.click(fn=run_all, inputs=inp, outputs=[out1, out2, out3])
 
 if __name__ == "__main__":
     demo.launch()
 
+