Update app.py

app.py

@@ -1,106 +1,76 @@
-import os
-import json
 import gradio as gr
 import time
 from transformers import pipeline
-from kaggle.api.kaggle_api_extended import KaggleApi
+from datasets import load_dataset
 
-# === Подготовка банковского набора данных через Kaggle ===
-# Пишем файл kaggle.json из переменных окружения, чтобы authenticate() нашёл его
-import json
-kaggle_config_dir = os.path.expanduser("~/.config/kaggle")
-os.makedirs(kaggle_config_dir, exist_ok=True)
-kaggle_json = {
-    "username": os.getenv("KAGGLE_USERNAME", ""),
-    "key": os.getenv("KAGGLE_KEY", "")
-}
-if not kaggle_json["username"] or not kaggle_json["key"]:
-    raise ValueError("Не найдены переменные окружения KAGGLE_USERNAME и KAGGLE_KEY")
-with open(os.path.join(kaggle_config_dir, "kaggle.json"), "w") as f:
-    json.dump(kaggle_json, f)
-
-# Скачиваем при первом запуске
-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
-
-if json_file is None:
-    raise FileNotFoundError('Не удалось найти JSON-файл с банковскими данными в ./data')
+# Загрузка датасета с отзывами банков
+dataset = load_dataset("Romjiik/Russian_bank_reviews", split="train")
 
-# === Инициализация трёх бесплатных русскоязычных моделей (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),
-}
+# Отбираем 2 примера для few-shot в CoT
+few_shot_examples = []
+for row in dataset.select(range(2)):
+    review = row["review"]
+    rating = row["rating"]
+    ex = f"Клиент: {review}\nОценка: {rating}\nОтвет: Пожалуйста, разъясните ситуацию в деталях. Мы поможем."
+    few_shot_examples.append(ex)
 
-# Системная инструкция для CoT
+# Системная инструкция
 system_instruction = (
-    "Вы — банковский ассистент. Ваша задача — корректно и вежливо отвечать на запросы клиентов банка, "
-    "давать рекомендации по банковским операциям и услугам."
+    "Ты — вежливый и точный банковский помощник. "
+    "Ты читаешь обращения клиентов и даешь корректные, подробные, официальные ответы. "
+    "Если данных недостаточно — просишь уточнение. Используй рассуждение шаг за шагом."
 )
 
-# Строим полный промпт с CoT и примерами
+# Загружаем три модели
+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),
+}
 
-def build_prompt(question: str) -> str:
-    few_shot_text = "\n\n".join(f"Пример:\n{ex}" for ex in examples)
-    prompt = (
+# Построение CoT-промпта
+def build_prompt(user_input):
+    examples = "\n\n".join(few_shot_examples)
+    return (
         f"{system_instruction}\n\n"
-        f"{few_shot_text}\n\n"
-        f"Вопрос клиента: {question}\n"
-        "Сначала подробно опишите рассуждения шаг за шагом, а затем кратко сформулируйте ответ."
+        f"{examples}\n\n"
+        f"Клиент: {user_input}\n"
+        f"Опиши шаг за шагом размышления, затем сформулируй окончательный ответ клиенту:"
     )
-    return prompt
-
-# Генерация ответов и измерение времени
 
-def generate(question: str):
-    prompt = build_prompt(question)
+# Генерация ответов
+def generate_answers(prompt):
     results = {}
     for name, pipe in models.items():
         start = time.time()
-        out = pipe(prompt, max_length=200, do_sample=True, top_p=0.9, temperature=0.7)[0]['generated_text']
+        out = pipe(prompt, max_length=300, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
         elapsed = round(time.time() - start, 2)
-        # Извлекаем связный ответ — последнюю строку
-        answer = out.strip().split('\n')[-1]
-        results[name] = {'answer': answer, 'time': elapsed}
+        final_line = out.strip().split('\n')[-1]
+        results[name] = {"answer": final_line, "time": elapsed}
     return results
 
-# Форматируем вывод для Gradio
-
-def format_outputs(question: str):
-    res = generate(question)
+# Формат вывода
+def run_models(user_input):
+    prompt = build_prompt(user_input)
+    res = generate_answers(prompt)
     return (
-        res['ruDialoGPT-small']['answer'], f"{res['ruDialoGPT-small']['time']}s",
-        res['ruDialoGPT-medium']['answer'], f"{res['ruDialoGPT-medium']['time']}s",
-        res['ruGPT3-small']['answer'], f"{res['ruGPT3-small']['time']}s"
+        res["ruDialoGPT-small"]["answer"], f"{res['ruDialoGPT-small']['time']} сек",
+        res["ruDialoGPT-medium"]["answer"], f"{res['ruDialoGPT-medium']['time']} сек",
+        res["ruGPT3-small"]["answer"], f"{res['ruGPT3-small']['time']} сек",
     )
 
-# === Интерфейс Gradio ===
+# Интерфейс Gradio
 with gr.Blocks() as demo:
-    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='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()
+    gr.Markdown("## 🤖 Банковский помощник: CoT + 3 модели (русский язык)")
+    inp = gr.Textbox(label="Запрос клиента", placeholder="Например: Я не могу попасть в личный кабинет", lines=2)
+    btn = gr.Button("Сгенерировать ответы")
+    out1 = gr.Textbox(label="ruDialoGPT-small")
+    t1 = gr.Textbox(label="Время")
+    out2 = gr.Textbox(label="ruDialoGPT-medium")
+    t2 = gr.Textbox(label="Время")
+    out3 = gr.Textbox(label="ruGPT3-small")
+    t3 = gr.Textbox(label="Время")
+
+    btn.click(run_models, inputs=[inp], outputs=[out1, t1, out2, t2, out3, t3])
 
+demo.launch()