Update app.py

app.py

@@ -1,53 +1,77 @@
 import gradio as gr
+import time
 from transformers import pipeline
+from datasets import load_dataset
 
-# Инициализация двух бесплатных русскоязычных моделей из Hugging Face
-# 1) ruDialoGPT-small — диалоговая модель на GPT-2 (MIT-лицензия) ([huggingface.co](https://huggingface.co/t-bank-ai/ruDialoGPT-small?utm_source=chatgpt.com))
-# 2) ruGPT3-small — семейство русскоязычных GPT-3 моделей от SberDevices (open-source) ([huggingface.co](https://huggingface.co/ai-forever/rugpt3small_based_on_gpt2?utm_source=chatgpt.com))
-
-pipe_dialo = pipeline(
-    task="text-generation",
-    model="t-bank-ai/ruDialoGPT-small",
-    tokenizer="t-bank-ai/ruDialoGPT-small",
-    device=-1  # CPU
-)
-
-pipe_rugpt3 = pipeline(
-    task="text-generation",
-    model="ai-forever/rugpt3small_based_on_gpt2",
-    tokenizer="ai-forever/rugpt3small_based_on_gpt2",
-    device=-1  # CPU
-)
-
-# Функция обработки пользовательского запроса
-# Возвращает генерацию от обеих моделей
-
-def generate_responses(prompt: str):
-    # Настройки генерации можно подкорректировать по потребностям
-    kwargs = {
-        "max_length": 200,
-        "do_sample": True,
-        "top_p": 0.9,
-        "temperature": 0.7
-    }
-    out1 = pipe_dialo(prompt, **kwargs)[0]["generated_text"]
-    out2 = pipe_rugpt3(prompt, **kwargs)[0]["generated_text"]
-    return out1, out2
-
-# Gradio-интерфейс
+# Загрузка бесплатных русскоязычных моделей
+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)
+}
 
+# Загрузка банковского датасета с диалогами
+bank_data = load_dataset('ai-lab/MBD', split='train')
+# Определяем колонку с диалогами
+col = next((c for c in bank_data.column_names if 'dialog' in c), None)
+if col is None:
+    raise ValueError('В датасете не найдена колонка с диалогами')
+# Берём первые два примера для few-shot
+examples = [item[col] for item in bank_data.select(range(2))]
+
+# Функция построения запроса с CoT и few-shot примерами
+
+def build_prompt(question):
+    few_shot = '\n\n'.join(f'Диалог:\n{ex}' for ex in examples)
+    prompt = (
+        f"{few_shot}\n\n"
+        f"Вопрос: {question}\n"
+        "Сначала подробно опишите рассуждения шаг за шагом, а затем дайте краткий связный ответ."
+    )
+    return prompt
+
+# Генерация ответов и снятие тайминга
+
+def generate(question):
+    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']
+        elapsed = round(time.time() - start, 2)
+        # Извлечение финального ответа после рассуждений
+        if 'Ответ:' in out:
+            answer = out.split('Ответ:')[-1].strip()
+        else:
+            answer = out.strip().split('\n')[-1]
+        results[name] = {'answer': answer, 'time': elapsed}
+    return results
+
+# Форматируем вывод для Gradio
+
+def format_outputs(question):
+    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"
+    )
+
+# Интерфейс Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## Русскоязычный чат с бесплатными моделями")
-    txt = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст на русском языке...", lines=2)
-    out1 = gr.Textbox(label="ruDialoGPT-small")
-    out2 = gr.Textbox(label="ruGPT3-small")
-    btn = gr.Button("Сгенерировать ответы")
-    btn.click(fn=generate_responses, inputs=txt, outputs=[out1, out2])
-
-# Запуск (для локального запуска и в Hugging Face Spaces)
-if __name__ == "__main__":
+    gr.Markdown('## CoT на трёх моделях и банковский датасет')
+    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()
 
 
 
 
+