Update app.py

app.py

@@ -1,42 +1,51 @@
 import gradio as gr
 import time
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 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
-    ),
-    '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
-    )
-}
+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
+)
 
-# Загрузка "мини" банковского датасета для few-shot (не более 10% данных)
+# 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
+)
+
+# Загрузка "мини" банковского датасета для few-shot (стриминг)
 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)
+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 примера, не загружая весь датасет
+# Собираем два few-shot примера
 examples = []
 for rec in bank_data_stream:
     examples.append(rec[col])
@@ -49,7 +58,7 @@ system_instruction = (
     " рассказывать о причинах и способах решения их проблем с банковскими услугами."
 )
 
-# Функция построения основного CoT-промпта
+# Построение CoT-промпта с few-shot
 
 def build_prompt(question: str) -> str:
     few_shot_text = '\n\n'.join(f"Пример диалога:\n{ex}" for ex in examples)
@@ -68,7 +77,7 @@ def generate(question: str):
     results = {}
     for name, pipe in models.items():
         start = time.time()
-        # Генерация по CoT
+        # для T5 используем text2text, для других text-generation
         out = pipe(
             prompt,
             max_length=400,
@@ -77,7 +86,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:
@@ -85,7 +94,7 @@ def generate(question: str):
         results[name] = {'answer': answer, 'time': elapsed}
     return results
 
-# Форматируем вывод для Gradio
+# Подготовка вывода для Gradio
 
 def format_outputs(question: str):
     res = generate(question)
@@ -98,20 +107,25 @@ def format_outputs(question: str):
 # Интерфейс Gradio
 with gr.Blocks() as demo:
     gr.Markdown('## Клиентские обращения: CoT на трёх моделях с MBD-mini и тайминг')
-    txt = gr.Textbox(label='Опишите проблему клиента', placeholder='Например: "Почему я не могу снять деньги с карты?"', lines=2)
+    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 Время')
-    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()
 
 
 
 
-
-