Update app.py

app.py

@@ -4,42 +4,54 @@ import time
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from datasets import load_dataset
 
+# 1) Публичные русскоязычные модели из RuGPT-3
 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)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# 3) Загрузка моделей и токенизаторов
 models = {}
-for label, name in MODEL_CONFIGS.items():
-    tokenizer = AutoTokenizer.from_pretrained(name)
-    model = AutoModelForCausalLM.from_pretrained(name)
-    model.to(device)
-    model.eval()
+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)
 
-# Загрузка датасета
+# 4) (По необходимости) загрузка датасета для примеров / дообучения
+#    Если не нужен — можно закомментировать
 load_dataset("ZhenDOS/alpha_bank_data", split="train")
 
-def cot_prompt_1(text):
-    return f"Клиент задал вопрос: {text}\nПодумай шаг за шагом и объясни, как бы ты ответил на это обращение от лица банка."
+# 5) CoT-промпты
+def cot_prompt_1(text: str) -> str:
+    return (
+        f"Клиент задал вопрос: «{text}»\n"
+        "Подумай шаг за шагом и подробно объясни ответ от лица банка."
+    )
 
-def cot_prompt_2(text):
-    return f"Вопрос клиента: {text}\nРазложи на части, что именно спрашивает клиент, и предложи логичный ответ с пояснениями."
+def cot_prompt_2(text: str) -> str:
+    return (
+        f"Вопрос клиента: «{text}»\n"
+        "Разложи на части, что именно спрашивает клиент, и предложи логичный ответ с пояснениями."
+    )
 
-def generate_all_responses(question):
+# 6) Генерация ответов и замер времени
+def generate_all_responses(question: str):
     results = {}
-    for model_name, (tokenizer, model) in models.items():
-        results[model_name] = {}
-        for i, prompt_func in enumerate([cot_prompt_1, cot_prompt_2], start=1):
-            prompt = prompt_func(question)
+    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.time()
+
+            start = time.time()
             with torch.no_grad():
-                outputs = model.generate(
+                output_ids = model.generate(
                     **inputs,
                     max_new_tokens=200,
                     do_sample=True,
@@ -47,36 +59,36 @@ def generate_all_responses(question):
                     top_p=0.9,
                     eos_token_id=tokenizer.eos_token_id
                 )
-            end_time = time.time()
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            response = response.replace(prompt, "").strip()
-            duration = round(end_time - start_time, 2)
-            results[model_name][f"CoT Промпт {i}"] = {
-                "response": response,
-                "time": f"{duration} сек."
+            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
 
-def display_responses(question):
-    all_responses = generate_all_responses(question)
-    output = ""
-    for model_name, prompts in all_responses.items():
-        output += f"\n### Модель: {model_name}\n"
-        for prompt_label, content in prompts.items():
-            output += f"\n**{prompt_label}** ({content['time']}):\n{content['response']}\n"
-    return output.strip()
+# 7) Оформление Markdown-вывода
+def display_responses(question: str) -> str:
+    all_res = generate_all_responses(question)
+    md = []
+    for model_name, prompts in all_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 — сравнение моделей",
-    description="Сравнение CoT-ответов от GigaChat, ChatGPT и DeepSeek-подобных моделей на обращение клиента.",
-    examples=[
-        "Как восстановить доступ в мобильный банк?",
-        "Почему с меня списали комиссию за обслуживание карты?",
-        "Какие условия по потребительскому кредиту?",
-    ]
+    inputs=gr.Textbox(lines=4, label="Введите вопрос клиента"),
+    outputs=gr.Markdown(label="Ответы трёх моделей"),
+    title="Alpha Bank Assistant — сравнение CoT-моделей",
+    description="Задайте вопрос клиентского обращения и сравните Chain-of-Thought ответы трёх русскоязычных моделей."
 )
 
 if __name__ == "__main__":