Update app.py

app.py

@@ -1,64 +1,73 @@
+import os
+import json
 import gradio as gr
 import time
 from transformers import pipeline
-from datasets import load_dataset
+from kaggle.api.kaggle_api_extended import KaggleApi
 
-# Инициализация трёх бесплатных русскоязычных моделей
-models = {
-    'ruDialoGPT-small': pipeline(
-        'text-generation',
-        model='t-bank-ai/ruDialoGPT-small',
-        tokenizer='t-bank-ai/ruDialoGPT-small',
-        device=-1
-    ),
-    'ruGPT3-small': pipeline(
-        'text-generation',
-        model='ai-forever/rugpt3small_based_on_gpt2',
-        tokenizer='ai-forever/rugpt3small_based_on_gpt2',
-        device=-1
-    ),
-    'rut5-small-chitchat': pipeline(
-        'text-generation',
-        model='cointegrated/rut5-small-chitchat',
-        tokenizer='cointegrated/rut5-small-chitchat',
-        device=-1
-    )
-}
+# === Подготовка банковского набора данных через Kaggle ===
+# Будет скачан dataset PromptCloudHQ/banking-chatbot-dataset, содержащий примеры вопросов и ответов для банковского чат-бота.
+DATA_DIR = './data'
+json_file = None
+# Скачиваем при первом запуске
+if not os.path.exists(DATA_DIR):
+    os.makedirs(DATA_DIR)
+    api = KaggleApi()
+    api.authenticate()
+    api.dataset_download_files('PromptCloudHQ/banking-chatbot-dataset', path=DATA_DIR, unzip=True)
+    # Находим JSON-файл с данными
+    for fname in os.listdir(DATA_DIR):
+        if fname.endswith('.json'):
+            json_file = os.path.join(DATA_DIR, fname)
+            break
+else:
+    # Если папка есть — ищем файл
+    for fname in os.listdir(DATA_DIR):
+        if fname.endswith('.json'):
+            json_file = os.path.join(DATA_DIR, fname)
+            break
 
-# Стриминг основного банковского датасета чтобы не загружать всё сразу
-bank_stream = load_dataset(
-    'ai-lab/MBD',
-    split='train',
-    streaming=True
-)
-# Используем явно колонку 'dialogs' из описания датасета
-# Берём первые два примера для few-shot
+if json_file is None:
+    raise FileNotFoundError('Не удалось найти JSON-файл с банковскими данными в ./data')
+
+# Загружаем JSON с примерами
+with open(json_file, 'r', encoding='utf-8') as f:
+    kb = json.load(f)
+# Структура: {"intents": [ ... ]}
+intents = kb.get('intents')
+if intents is None:
+    raise ValueError('Ожидался ключ "intents" в JSON-файле датасета')
+
+# Собираем два few-shot примера
 examples = []
-for record in bank_stream:
-    if 'dialogs' in record:
-        examples.append(record['dialogs'])
-    elif 'dialog_embeddings' in record:
-        examples.append(record['dialog_embeddings'])
-    if len(examples) == 2:
-        break
-if len(examples) < 2:
-    raise ValueError('Не удалось получить два примера dialog из MBD')
+for intent in intents[:2]:
+    patterns = intent.get('patterns', [])
+    responses = intent.get('responses', [])
+    ex = f"Паттерны: {', '.join(patterns)}\nОтветы: {', '.join(responses)}"
+    examples.append(ex)
+
+# === Инициализация трёх бесплатных русскоязычных моделей (GPT-2 based) ===
+models = {
+    'ruDialoGPT-small': pipeline('text-generation', model='t-bank-ai/ruDialoGPT-small', tokenizer='t-bank-ai/ruDialoGPT-small', device=-1),
+    'ruDialoGPT-medium': pipeline('text-generation', model='t-bank-ai/ruDialoGPT-medium', tokenizer='t-bank-ai/ruDialoGPT-medium', device=-1),
+    'ruGPT3-small': pipeline('text-generation', model='ai-forever/rugpt3small_based_on_gpt2', tokenizer='ai-forever/rugpt3small_based_on_gpt2', device=-1),
+}
 
 # Системная инструкция для CoT
 system_instruction = (
-    "Вы — банковский ассистент. Ваша задача — корректно и вежливо отвечать на запросы клиентов банка,"
-    " рассказывать о причинах и способах решения их проблем с банковскими услугами."
+    "Вы — банковский ассистент. Ваша задача — корректно и вежливо отвечать на запросы клиентов банка, "
+    "давать рекомендации по банковским операциям и услугам."
 )
 
-# Функция построения CoT промпта с few-shot примерами
+# Строим полный промпт с CoT и примерами
 
 def build_prompt(question: str) -> str:
-    few_shot = '\n\n'.join(f"Пример диалога:\n{ex}" for ex in examples)
+    few_shot_text = "\n\n".join(f"Пример:\n{ex}" for ex in examples)
     prompt = (
         f"{system_instruction}\n\n"
-        f"{few_shot}\n\n"
+        f"{few_shot_text}\n\n"
         f"Вопрос клиента: {question}\n"
-        "Сначала подробно опишите рассуждения шаг за шагом, а затем дайте краткий связный ответ."
+        "Сначала подробно опишите рассуждения шаг за шагом, а затем кратко сформулируйте ответ."
     )
     return prompt
 
@@ -69,47 +78,33 @@ def generate(question: str):
     results = {}
     for name, pipe in models.items():
         start = time.time()
-        out = pipe(
-            prompt,
-            max_length=400,
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.7
-        )[0]['generated_text']
+        out = pipe(prompt, max_length=200, do_sample=True, top_p=0.9, temperature=0.7)[0]['generated_text']
         elapsed = round(time.time() - start, 2)
-        # Извлечение итогового ответа после 'Ответ:' или последней строки
-        if 'Ответ:' in out:
-            answer = out.split('Ответ:')[-1].strip()
-        else:
-            answer = out.strip().split('\n')[-1]
+        # Извлекаем связный ответ — последнюю строку
+        answer = out.strip().split('\n')[-1]
         results[name] = {'answer': answer, 'time': elapsed}
     return results
 
-# Форматируем данные для Gradio интерфейса
+# Форматируем вывод для Gradio
 
 def format_outputs(question: str):
     res = generate(question)
     return (
         res['ruDialoGPT-small']['answer'], f"{res['ruDialoGPT-small']['time']}s",
-        res['ruGPT3-small']['answer'], f"{res['ruGPT3-small']['time']}s",
-        res['rut5-small-chitchat']['answer'], f"{res['rut5-small-chitchat']['time']}s"
+        res['ruDialoGPT-medium']['answer'], f"{res['ruDialoGPT-medium']['time']}s",
+        res['ruGPT3-small']['answer'], f"{res['ruGPT3-small']['time']}s"
     )
 
-# Графический интерфейс Gradio
-
+# === Интерфейс Gradio ===
 with gr.Blocks() as demo:
-    gr.Markdown('## Ответы на клиентские обращения с CoT на трёх моделях и таймингом')
-    txt = gr.Textbox(
-        label='Описание проблемы клиента',
-        placeholder='Например: "Почему я не могу снять деньги с карты?"',
-        lines=2
-    )
-    btn = gr.Button('Сгенерировать ответ')
+    gr.Markdown("## Ответы на клиентские обращения\nCoT + тайминг по трём бесплатным моделям")
+    txt = gr.Textbox(label='Описание проблемы клиента', placeholder='Например: "Почему я не могу снять деньги с карты?"', lines=2)
+    btn = gr.Button('Сгенерировать ответы')
     out1 = gr.Textbox(label='ruDialoGPT-small Ответ')
     t1 = gr.Textbox(label='ruDialoGPT-small Время')
-    out2 = gr.Textbox(label='ruGPT3-small Ответ')
-    t2 = gr.Textbox(label='ruGPT3-small Время')
-    out3 = gr.Textbox(label='rut5-small-chitchat Ответ')
-    t3 = gr.Textbox(label='rut5-small-chitchat Время')
+    out2 = gr.Textbox(label='ruDialoGPT-medium Ответ')
+    t2 = gr.Textbox(label='ruDialoGPT-medium Время')
+    out3 = gr.Textbox(label='ruGPT3-small Ответ')
+    t3 = gr.Textbox(label='ruGPT3-small Время')
     btn.click(format_outputs, inputs=[txt], outputs=[out1, t1, out2, t2, out3, t3])
     demo.launch()