Update app.py

app.py

@@ -1,56 +1,48 @@
 import gradio as gr
 import time
-from transformers import pipeline, AutoTokenizer
+from transformers import pipeline
 from datasets import load_dataset
 
 # Инициализация трёх бесплатных русскоязычных моделей
-models = {}
-
-# 1) ruDialoGPT-small
-models['ruDialoGPT-small'] = pipeline(
-    'text-generation',
-    model='t-bank-ai/ruDialoGPT-small',
-    tokenizer='t-bank-ai/ruDialoGPT-small',
-    device=-1
-)
-
-# 2) ruGPT3-small
-models['ruGPT3-small'] = pipeline(
-    'text-generation',
-    model='ai-forever/rugpt3small_based_on_gpt2',
-    tokenizer='ai-forever/rugpt3small_based_on_gpt2',
-    device=-1
-)
-
-# 3) rut5-small-chitchat (T5 requires text2text и slow tokenizer)
-t5_tokenizer = AutoTokenizer.from_pretrained(
-    'cointegrated/rut5-small-chitchat',
-    use_fast=False
-)
-models['rut5-small-chitchat'] = pipeline(
-    'text2text-generation',
-    model='cointegrated/rut5-small-chitchat',
-    tokenizer=t5_tokenizer,
-    device=-1
-)
+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
+    )
+}
 
-# Загрузка "мини" банковского датасета для few-shot (стриминг)
-bank_data_stream = load_dataset(
-    'ai-lab/MBD-mini',
+# Стриминг основного банковского датасета чтобы не загружать всё сразу
+bank_stream = load_dataset(
+    'ai-lab/MBD',
     split='train',
     streaming=True
 )
-# Определяем колонку с диалогами по ключам
-first_record = next(iter(bank_data_stream))
-col = next((c for c in first_record.keys() if 'dialog' in c.lower() or 'диалог' in c.lower()), None)
-if col is None:
-    raise ValueError('Не найдена колонка с диалогами в MBD-mini')
-# Собираем два few-shot примера
+# Используем явно колонку 'dialogs' из описания датасета
+# Берём первые два примера для few-shot
 examples = []
-for rec in bank_data_stream:
-    examples.append(rec[col])
+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')
 
 # Системная инструкция для CoT
 system_instruction = (
@@ -58,13 +50,13 @@ system_instruction = (
     " рассказывать о причинах и способах решения их проблем с банковскими услугами."
 )
 
-# Построение CoT-промпта с few-shot
+# Функция построения CoT промпта с few-shot примерами
 
 def build_prompt(question: str) -> str:
-    few_shot_text = '\n\n'.join(f"Пример диалога:\n{ex}" for ex in examples)
+    few_shot = '\n\n'.join(f"Пример диалога:\n{ex}" for ex in examples)
     prompt = (
         f"{system_instruction}\n\n"
-        f"{few_shot_text}\n\n"
+        f"{few_shot}\n\n"
         f"Вопрос клиента: {question}\n"
         "Сначала подробно опишите рассуждения шаг за шагом, а затем дайте краткий связный ответ."
     )
@@ -77,7 +69,6 @@ def generate(question: str):
     results = {}
     for name, pipe in models.items():
         start = time.time()
-        # для T5 используем text2text, для других text-generation
         out = pipe(
             prompt,
             max_length=400,
@@ -86,7 +77,7 @@ def generate(question: str):
             temperature=0.7
         )[0]['generated_text']
         elapsed = round(time.time() - start, 2)
-        # Извлечь финальный ответ
+        # Извлечение итогового ответа после 'Ответ:' или последней строки
         if 'Ответ:' in out:
             answer = out.split('Ответ:')[-1].strip()
         else:
@@ -94,7 +85,7 @@ def generate(question: str):
         results[name] = {'answer': answer, 'time': elapsed}
     return results
 
-# Подготовка вывода для Gradio
+# Форматируем данные для Gradio интерфейса
 
 def format_outputs(question: str):
     res = generate(question)
@@ -104,11 +95,12 @@ def format_outputs(question: str):
         res['rut5-small-chitchat']['answer'], f"{res['rut5-small-chitchat']['time']}s"
     )
 
-# Интерфейс Gradio
+# Графический интерфейс Gradio
+
 with gr.Blocks() as demo:
-    gr.Markdown('## Клиентские обращения: CoT на трёх моделях с MBD-mini и тайминг')
+    gr.Markdown('## Ответы на клиентские обращения с CoT на трёх моделях и таймингом')
     txt = gr.Textbox(
-        label='Опишите проблему клиента',
+        label='Описание проблемы клиента',
         placeholder='Например: "Почему я не могу снять деньги с карты?"',
         lines=2
     )
@@ -119,13 +111,5 @@ with gr.Blocks() as demo:
     t2 = gr.Textbox(label='ruGPT3-small Время')
     out3 = gr.Textbox(label='rut5-small-chitchat Ответ')
     t3 = gr.Textbox(label='rut5-small-chitchat Время')
-    btn.click(
-        format_outputs,
-        inputs=[txt],
-        outputs=[out1, t1, out2, t2, out3, t3]
-    )
+    btn.click(format_outputs, inputs=[txt], outputs=[out1, t1, out2, t2, out3, t3])
     demo.launch()
-
-
-
-