Update app.py

app.py

@@ -1,144 +1,83 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
+import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from datasets import load_dataset
 
-# Загрузка датасета
-dataset = load_dataset("ZhenDOS/alpha_bank_data")
+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"
+}
 
-# Инициализация разных моделей
-model_name = "ai-forever/rugpt3large_based_on_gpt2"
-    "GigaChat-like:" "ai-forever/rugpt3large_based_on_gpt2",  # Русская модель большого размера
-    "ChatGPT-like": "tinkoff-ai/ruDialoGPT-medium",           # Диалоговая модель для русского языка
-    "DeepSeek-like": "ai-forever/sbert_large_nlu_ru"           # Русская модель для понимания текста
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+models = {}
+for label, name in MODEL_CONFIGS.items():
+    tokenizer = AutoTokenizer.from_pretrained(name)
+    model = AutoModelForCausalLM.from_pretrained(name)
+    model.to(device)
+    model.eval()
+    models[label] = (tokenizer, model)
 
-# Инициализация моделей и токенизаторов
-models = model_name
-tokenizers = AutoTokenizer.from_pretrained(model_name)
+# Загрузка датасета
+load_dataset("ZhenDOS/alpha_bank_data", split="train")
 
-for model_name, model_path in MODELS.items():
-    try:
-        if model_name == "DeepSeek-like":
-            # Для SBERT используем pipeline
-            models[model_name] = pipeline("text-generation", model=model_path)
-        else:
-            tokenizers[model_name] = AutoTokenizer.from_pretrained(model_path)
-            models[model_name] = AutoModelForCausalLM.from_pretrained(model_path)
-    except Exception as e:
-        print(f"Ошибка при загрузке модели {model_name}: {e}")
+def cot_prompt_1(text):
+    return f"Клиент задал вопрос: {text}\nПодумай шаг за шагом и объясни, как бы ты ответил на это обращение от лица банка."
 
-# Промпты для обработки обращений
-PROMPTS = {
-    "Анализ проблемы": 
-        "Проанализируй клиентское обращение и выдели основную проблему. "
-        "Обращение: {text}\n\nПроблема:",
-    
-    "Формирование ответа":
-        "Клиент обратился с проблемой: {problem}\n\n"
-        "Сформируй вежливый и профессиональный ответ, предлагая решение. "
-        "Используй информацию о банковских услугах. Ответ:"
-}
+def cot_prompt_2(text):
+    return f"Вопрос клиента: {text}\nРазложи на части, что именно спрашивает клиент, и предложи логичный ответ с пояснениями."
 
-def generate_with_model(prompt, model_name, max_length=150):
-    """Генерация ответа с помощью выбранной модели"""
-    if model_name not in models:
-        return f"Модель {model_name} не загружена"
-    
-    try:
-        if model_name == "DeepSeek-like":
-            # Обработка через pipeline
-            result = models[model_name](
-                prompt,
-                max_length=max_length,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9
-            )
-            return result[0]['generated_text']
-        else:
-            # Обработка через transformers
-            inputs = tokenizers[model_name](prompt, return_tensors="pt", truncation=True)
-            
+def generate_all_responses(question):
+    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)
+            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()
             with torch.no_grad():
-                outputs = models[model_name].generate(
+                outputs = model.generate(
                     **inputs,
-                    max_new_tokens=max_length,
+                    max_new_tokens=200,
                     do_sample=True,
                     temperature=0.7,
                     top_p=0.9,
-                    eos_token_id=tokenizers[model_name].eos_token_id
+                    eos_token_id=tokenizer.eos_token_id
                 )
-            
-            response = tokenizers[model_name].decode(outputs[0], skip_special_tokens=True)
-            return response[len(prompt):] if response.startswith(prompt) else response
-    except Exception as e:
-        return f"Ошибка генерации: {str(e)}"
-
-def process_complaint(text, prompt_type):
-    """Обработка клиентского обращения с выбранным промптом"""
-    if prompt_type not in PROMPTS:
-        return "Неверный тип промпта"
-    
-    # Получаем случайный пример из датасета, если текст не введен
-    if not text.strip():
-        example = dataset['train'].shuffle().select(range(1))[0]
-        text = example['text']
-    
-    prompt = PROMPTS[prompt_type].format(text=text, problem="")
-    
-    results = {}
-    for model_name in MODELS.keys():
-        results[model_name] = generate_with_model(prompt, model_name)
-    
+            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} сек."
+            }
     return results
 
-# Интерфейс Gradio
-with gr.Blocks(title="Анализ клиентских обращений Alpha Bank") as demo:
-    gr.Markdown("## Анализ клиентских обращений Alpha Bank")
-    gr.Markdown("Тестирование разных моделей на обработку обращений")
-    
-    with gr.Row():
-        with gr.Column():
-            text_input = gr.Textbox(
-                label="Текст обращения",
-                placeholder="Введите текст обращения или оставьте пустым для примера из датасета",
-                lines=5
-            )
-            prompt_type = gr.Radio(
-                list(PROMPTS.keys()),
-                label="Тип промпта",
-                value=list(PROMPTS.keys())[0]
-            )
-            submit_btn = gr.Button("Обработать")
-        
-        with gr.Column():
-            outputs = []
-            for model_name in MODELS.keys():
-                outputs.append(
-                    gr.Textbox(
-                        label=f"{model_name}",
-                        interactive=False,
-                        lines=5
-                    )
-                )
-    
-    # Примеры из датасета
-    examples = gr.Examples(
-        examples=[x['text'] for x in dataset['train'].select(range(3))],
-        inputs=text_input,
-        label="Примеры из датасета"
-    )
-    
-    def process_and_display(text, prompt_type):
-        results = process_complaint(text, prompt_type)
-        return [results.get(model_name, "") for model_name in MODELS.keys()]
-    
-    submit_btn.click(
-        fn=process_and_display,
-        inputs=[text_input, prompt_type],
-        outputs=outputs
-    )
+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()
+
+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=[
+        "Как восстановить доступ в мобильный банк?",
+        "Почему с меня списали комиссию за обслуживание карты?",
+        "Какие условия по потребительскому кредиту?",
+    ]
+)
 
 if __name__ == "__main__":
     demo.launch()