Update app.py

app.py

@@ -1,100 +1,144 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 from datasets import load_dataset
 
-# Загружаем модель и токенизатор
-model_name = "ai-forever/rugpt3small_based_on_gpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Загрузка датасета
+dataset = load_dataset("ZhenDOS/alpha_bank_data")
 
-# Дополнительные знания о банковских услугах
-BANK_KNOWLEDGE = {
-    "Как проверить баланс карты?": [
-        "1. Через мобильное приложение банка (раздел 'Карты' → 'Баланс')",
-        "2. В интернет-банке (в личном кабинете выберите карту)",
-        "3. По SMS (отправьте BALANCE на номер 900)",
-        "4. В банкомате (вставьте карту и выберите 'Запрос баланса')",
-        "5. По телефону горячей линии (8-800-100-00-00)"
-    ],
-    "Как восстановить утерянную карту?": [
-        "1. Немедленно позвоните в банк по телефону 8-800-100-00-00 для блокировки карты",
-        "2. Обратитесь в отделение банка с паспортом",
-        "3. Заполните заявление на перевыпуск карты",
-        "4. Новая карта будет готова через 3-5 рабочих дней"
-    ]
+# Инициализация разных моделей
+MODELS = {
+    "GigaChat-like": "ai-forever/rugpt3large_based_on_gpt2",  # Русская модель большого размера
+    "ChatGPT-like": "tinkoff-ai/ruDialoGPT-medium",           # Диалоговая модель для русского языка
+    "DeepSeek-like": "ai-forever/sbert_large_nlu_ru"           # Русская модель для понимания текста
 }
 
-def enhance_response(question, generated_response):
-    # Если вопрос есть в наших знаниях, возвращаем структурированный ответ
-    if question in BANK_KNOWLEDGE:
-        return "\n".join(BANK_KNOWLEDGE[question])
-    
-    # Улучшаем стандартные ответы модели
-    improvements = {
-        "баланс": "Вы можете проверить баланс карты:\n"
-                 "1. В мобильном приложении\n"
-                 "2. Через интернет-банк\n"
-                 "3. В банкомате\n"
-                 "4. По телефону горячей линии 8-800-100-00-00",
-        "кредит": "По вопросам кредитования вы можете:\n"
-                 "1. Оставить заявку на сайте\n"
-                 "2. Обратиться в отделение банка\n"
-                 "3. Позвонить по телефону 8-800-100-00-00",
-        "карт": "По вопросам банковских карт:\n"
-               "1. Обратитесь в отделение банка\n"
-               "2. Позвоните на горячую линию\n"
-               "3. Используйте чат в мобильном приложении"
-    }
-    
-    for keyword, improved_answer in improvements.items():
-        if keyword in question.lower():
-            return improved_answer
+# Инициализация моделей и токенизаторов
+models = {}
+tokenizers = {}
+
+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}")
+
+# Промпты для обработки обращений
+PROMPTS = {
+    "Анализ проблемы": 
+        "Проанализируй клиентское обращение и выдели основную проблему. "
+        "Обращение: {text}\n\nПроблема:",
     
-    return generated_response
+    "Формирование ответа":
+        "Клиент обратился с проблемой: {problem}\n\n"
+        "Сформируй вежливый и профессиональный ответ, предлагая решение. "
+        "Используй информацию о банковских услугах. Ответ:"
+}
 
-def generate_response(prompt):
-    # Генерируем ответ с помощью модели
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+def generate_with_model(prompt, model_name, max_length=150):
+    """Генерация ответа с помощью выбранной модели"""
+    if model_name not in models:
+        return f"Модель {model_name} не загружена"
     
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=150,
-            do_sample=True,
-            temperature=0.7,
-            top_k=50,
-            top_p=0.95,
-            eos_token_id=tokenizer.eos_token_id,
-            no_repeat_ngram_size=3,
-            early_stopping=True
-        )
+    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)
+            
+            with torch.no_grad():
+                outputs = models[model_name].generate(
+                    **inputs,
+                    max_new_tokens=max_length,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    eos_token_id=tokenizers[model_name].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 "Неверный тип промпта"
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Получаем случайный пример из датасета, если текст не введен
+    if not text.strip():
+        example = dataset['train'].shuffle().select(range(1))[0]
+        text = example['text']
     
-    # Удаляем промпт из ответа
-    if response.startswith(prompt):
-        response = response[len(prompt):].strip()
+    prompt = PROMPTS[prompt_type].format(text=text, problem="")
     
-    # Улучшаем ответ
-    enhanced_response = enhance_response(prompt, response)
+    results = {}
+    for model_name in MODELS.keys():
+        results[model_name] = generate_with_model(prompt, model_name)
     
-    return enhanced_response
+    return results
 
 # Интерфейс Gradio
-demo = gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(lines=4, label="Введите вопрос по клиентским обращениям в банк"),
-    outputs=gr.Textbox(label="Ответ модели"),
-    title="Анализ клиентских обращений — Alpha Bank Assistant",
-    description="Получите точные ответы на вопросы о банковских услугах",
-    examples=[
-        "Как проверить баланс карты?",
-        "Как восстановить утерянную карту?",
-        "Как оформить кредитную карту?",
-        "Какие документы нужны для открытия счета?"
-    ]
-)
+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
+    )
 
 if __name__ == "__main__":
     demo.launch()