Update app.py

app.py

@@ -3,43 +3,81 @@ import time
 from transformers import pipeline
 from datasets import load_dataset
 
-# Загрузка бесплатных русскоязычных моделей
+# Инициализация трёх бесплатных русскоязычных моделей
 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)
+    '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
+    )
 }
 
-# Загрузка банковского датасета с диалогами
-bank_data = load_dataset('ai-lab/MBD', split='train')
-# Определяем колонку с диалогами
-col = next((c for c in bank_data.column_names if 'dialog' in c), None)
+# Загрузка "мини" банковского датасета для few-shot (не более 10% данных)
+bank_data_stream = load_dataset(
+    'ai-lab/MBD-mini',
+    split='train',
+    streaming=True
+)
+# Определяем колонку с диалогами среди ключей
+first_record = next(iter(bank_data_stream))
+col = next((c for c in first_record.keys() if 'dialog' in c or 'Диалог' in c or 'dialogue' in c.lower()), None)
 if col is None:
-    raise ValueError('В датасете не найдена колонка с диалогами')
-# Берём первые два примера для few-shot
-examples = [item[col] for item in bank_data.select(range(2))]
+    raise ValueError('Не найдена колонка с диалогами в MBD-mini')
+# Собираем два few-shot примера, не загружая весь датасет
+examples = []
+for rec in bank_data_stream:
+    examples.append(rec[col])
+    if len(examples) == 2:
+        break
+
+# Системная инструкция для CoT
+system_instruction = (
+    "Вы — банковский ассистент. Ваша задача — корректно и вежливо отвечать на запросы клиентов банка,"
+    " рассказывать о причинах и способах решения их проблем с банковскими услугами."
+)
 
-# Функция построения запроса с CoT и few-shot примерами
+# Функция построения основного CoT-промпта
 
-def build_prompt(question):
-    few_shot = '\n\n'.join(f'Диалог:\n{ex}' for ex in examples)
+def build_prompt(question: str) -> str:
+    few_shot_text = '\n\n'.join(f"Пример диалога:\n{ex}" for ex in examples)
     prompt = (
-        f"{few_shot}\n\n"
-        f"Вопрос: {question}\n"
+        f"{system_instruction}\n\n"
+        f"{few_shot_text}\n\n"
+        f"Вопрос клиента: {question}\n"
         "Сначала подробно опишите рассуждения шаг за шагом, а затем дайте краткий связный ответ."
     )
     return prompt
 
-# Генерация ответов и снятие тайминга
+# Генерация ответов и измерение времени
 
-def generate(question):
+def generate(question: str):
     prompt = build_prompt(question)
     results = {}
     for name, pipe in models.items():
         start = time.time()
-        out = pipe(prompt, max_length=300, do_sample=True, top_p=0.9, temperature=0.7)[0]['generated_text']
+        # Генерация по CoT
+        out = pipe(
+            prompt,
+            max_length=400,
+            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:
@@ -49,25 +87,26 @@ def generate(question):
 
 # Форматируем вывод для Gradio
 
-def format_outputs(question):
+def format_outputs(question: str):
     res = generate(question)
     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['ruGPT3-small']['answer'], f"{res['ruGPT3-small']['time']}s",
+        res['rut5-small-chitchat']['answer'], f"{res['rut5-small-chitchat']['time']}s"
     )
 
 # Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown('## CoT на трёх моделях и банковский датасет')
-    txt = gr.Textbox(label='Ваш вопрос', placeholder='Введите вопрос, связанный с банковскими услугами', lines=2)
-    btn = gr.Button('Сгенерировать')
+    gr.Markdown('## Клиентские обращения: CoT на трёх моделях с MBD-mini и тайминг')
+    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 Время')
+    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 Время')
     btn.click(format_outputs, inputs=[txt], outputs=[out1, t1, out2, t2, out3, t3])
     demo.launch()
 
@@ -75,3 +114,4 @@ with gr.Blocks() as demo:
 
 
 
+