Update app.py

app.py

@@ -2,94 +2,75 @@ import gradio as gr
 import torch
 import time
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from datasets import load_dataset
 
-# 1) Публичные русскоязычные модели из RuGPT-3
+# 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) Устройство (GPU если есть, иначе CPU)
+# 2) Выбор устройства
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 3) Загрузка моделей и токенизаторов
+# 3) Загрузка токенизаторов и моделей один раз при старте
 models = {}
 for label, repo_id in MODEL_CONFIGS.items():
-    tokenizer = AutoTokenizer.from_pretrained(repo_id)
-    model     = AutoModelForCausalLM.from_pretrained(repo_id)
-    model.to(device).eval()
-    models[label] = (tokenizer, model)
+    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}")
 
-# 4) (По необходимости) загрузка датасета для примеров / дообучения
-#    Если не нужен — можно закомментировать
-load_dataset("ZhenDOS/alpha_bank_data", split="train")
+# 4) Chain-of-Thought промпты
+def cot_prompt_1(q): return f"Клиент: «{q}»\nШаг за шагом объясни ответ от лица банка."
+def cot_prompt_2(q): return f"Клиент: «{q}»\nРазбери вопрос на части и дай развёрнутый ответ."
 
-# 5) CoT-промпты
-def cot_prompt_1(text: str) -> str:
-    return (
-        f"Клиент задал вопрос: «{text}»\n"
-        "Подумай шаг за шагом и подробно объясни ответ от лица банка."
-    )
+# 5) Функция генерации
+def generate_all_responses(question):
+    if not question.strip():
+        return {k: {"error": "Пустой вопрос"} for k in models}
+    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": "-"}
+    return out
 
-def cot_prompt_2(text: str) -> str:
-    return (
-        f"Вопрос клиента: «{text}»\n"
-        "Разложи на части, что именно спрашивает клиент, и предложи логичный ответ с пояснениями."
-    )
-
-# 6) Генерация ответов и замер времени
-def generate_all_responses(question: str):
-    results = {}
-    for name, (tokenizer, model) in models.items():
-        results[name] = {}
-        for idx, prompt_fn in enumerate([cot_prompt_1, cot_prompt_2], start=1):
-            prompt = prompt_fn(question)
-            inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-            inputs = {k: v.to(device) for k, v in inputs.items()}
-
-            start = time.time()
-            with torch.no_grad():
-                output_ids = model.generate(
-                    **inputs,
-                    max_new_tokens=200,
-                    do_sample=True,
-                    temperature=0.7,
-                    top_p=0.9,
-                    eos_token_id=tokenizer.eos_token_id
-                )
-            latency = round(time.time() - start, 2)
-
-            text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-            # Убираем повтор промпта
-            if text.startswith(prompt):
-                text = text[len(prompt):].strip()
-
-            results[name][f"CoT-промпт {idx}"] = {
-                "response": text,
-                "time":     f"{latency} сек."
-            }
-    return results
-
-# 7) Оформление Markdown-вывода
-def display_responses(question: str) -> str:
-    all_res = generate_all_responses(question)
+# 6) Обёртка для Gradio
+def run_all(question):
+    res = generate_all_responses(question)
     md = []
-    for model_name, prompts in all_res.items():
-        md.append(f"## Модель: **{model_name}**")
+    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)
 
-# 8) Интерфейс Gradio
-demo = gr.Interface(
-    fn=display_responses,
-    inputs=gr.Textbox(lines=4, label="Введите вопрос клиента"),
-    outputs=gr.Markdown(label="Ответы трёх моделей"),
-    title="Alpha Bank Assistant — сравнение CoT-моделей",
-    description="Задайте вопрос клиентского обращения и сравните Chain-of-Thought ответы трёх русскоязычных моделей."
-)
+# 7) Интерфейс Gradio с блоками
+with gr.Blocks() as demo:
+    gr.Markdown("# Alpha Bank Assistant — сравнение CoT-моделей")
+    inp = gr.Textbox(lines=3, placeholder="Введите вопрос клиента...", label="Вопрос клиента")
+    btn = gr.Button("Сгенерировать ответы")
+    out = gr.Markdown(label="Результаты")
+    btn.click(fn=run_all, inputs=inp, outputs=out)
 
 if __name__ == "__main__":
     demo.launch()
+