Update app.py

app.py

@@ -1,15 +1,33 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
+from datasets import load_dataset
 
-# Загружаем модель и токенизатор
+# Загружаем модель, токенизатор и датасет
 model_name = "ai-forever/rugpt3small_based_on_gpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Функция генерации ответа
+# Загружаем банковский датасет для контекста
+bank_dataset = load_dataset("ZhenDOS/alpha_bank_data")
+
+# Создаем контекст из датасета (первые несколько примеров)
+context_examples = "\n".join([
+    f"Вопрос: {example['question']}\nОтвет: {example['answer']}" 
+    for example in bank_dataset['train'].select(range(5))
+])
+
+# Функция генерации ответа с учетом банковского контекста
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    # Добавляем контекст из датасета к промпту
+    full_prompt = f"""Контекст по банковским вопросам:
+{context_examples}
+
+Вопрос клиента: {prompt}
+Ответ:"""
+    
+    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512)
+    
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -18,21 +36,39 @@ def generate_response(prompt):
             temperature=0.7,
             top_k=50,
             top_p=0.95,
-            eos_token_id=tokenizer.eos_token_id
+            eos_token_id=tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
+            early_stopping=True
         )
+    
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Удаляем промпт из ответа, если повторяется
-    if response.startswith(prompt):
-        response = response[len(prompt):].strip()
-    return response.strip()
+    
+    # Удаляем промпт из ответа
+    if response.startswith(full_prompt):
+        response = response[len(full_prompt):].strip()
+    
+    # Постобработка ответа
+    response = response.split("\n")[0]  # Берем только первую строку ответа
+    response = response.replace("Ответ:", "").strip()
+    
+    return response
+
+# Интерфейс Gradio с примерами вопросов
+examples = [
+    "Как восстановить утерянную карту?",
+    "Какие документы нужны для открытия счета?",
+    "Как проверить баланс карты?",
+    "Как оформить кредитную карту?",
+    "Какие комиссии за перевод между счетами?"
+]
 
-# Интерфейс Gradio
 demo = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(lines=4, label="Введите вопрос по клиентским обещаниям банка"),
+    inputs=gr.Textbox(lines=4, label="Введите вопрос по клиентским обращениям в банк"),
     outputs=gr.Textbox(label="Ответ модели"),
-    title="Анализ клиентских обещаний — RuGPT-3",
-    description="Используется модель ai-forever/rugpt3small_based_on_gpt2 на основе данных ZhenDOS/alpha_bank_data."
+    title="Анализ клиентских обращений — RuGPT-3 с Alpha Bank Data",
+    description="Используется модель ai-forever/rugpt3small_based_on_gpt2, дообученная на данных ZhenDOS/alpha_bank_data.",
+    examples=examples
 )
 
 # Запуск