Update app.py

app.py

@@ -1,61 +1,46 @@
 import gradio as gr
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from datasets import load_dataset
-
-# Загрузим модель и токенизатор T5
-model_name = "t5-small"  # Или используйте более крупную модель для лучшего качества
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-
-# Загрузим датасет Banking77
-dataset = load_dataset("banking77")
-
-# Подготовим данные
-# Для этой задачи будем использовать только столбцы вопроса и ответа
-train_data = dataset['train']
-test_data = dataset['test']
-
-# Пример одного вопроса и ответа из набора данных
-example = train_data[0]
-print(f"Вопрос: {example['text']}")
-print(f"Ответ: {example['label']}")
-
-# Функция генерации ответа на запрос
-def respond_to_client(message):
-    # Формируем промпт для генерации ответа
-    prompt = f"Вопрос: {message} Ответ банка:"
-    
-    # Токенизируем промпт
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-
-    # Генерируем ответ
-    output = model.generate(
-        inputs["input_ids"], 
-        max_length=200,  # Максимальная длина ответа
-        num_return_sequences=3,  # Количество ответов
-        no_repeat_ngram_size=2,  # Предотвращение повторов
-        top_p=0.95,  # Топ P для сэмплинга
-        top_k=50,  # Топ K для сэмплинга
-        temperature=0.7,  # Температура (чем ниже, тем более детерминированный ответ)
-        do_sample=True  # Использование сэмплинга для разнообразных ответов
-    )
-
-    # Декодируем ответы
-    responses = [tokenizer.decode(output[i], skip_special_tokens=True) for i in range(3)]
-    
-    return "\n\n".join(responses)
-
-# Создаем интерфейс с помощью Gradio
+import random
+import torch
+
+# Загружаем датасет banking77
+dataset = load_dataset("banking77", split="train")
+
+# Пример перевода некоторых вопросов вручную
+translated_questions = {
+    "How do I activate my debit card?": "Как активировать мою дебетовую карту?",
+    "What are the fees for international transfers?": "Какие комиссии за международные переводы?",
+    "How do I reset my password?": "Как сбросить мой пароль?",
+    "Where can I find my IBAN number?": "Где мне найти мой номер IBAN?",
+    "How to close my bank account?": "Как закрыть мой банковский счет?",
+}
+
+# Загружаем русскую модель
+model_name = "cointegrated/rugpt2-large"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Генерация ответа
+def generate_response(question):
+    prompt = f"Клиент спрашивает: {question}\nБанк отвечает:"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True, top_p=0.95, top_k=50)
+    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = generated.replace(prompt, "").strip()
+    return response
+
+# Интерфейс
 iface = gr.Interface(
-    fn=respond_to_client,  # Функция обработки запроса
-    inputs=gr.Textbox(lines=2, placeholder="Введите вопрос клиента..."),  # Ввод
-    outputs="text",  # Вывод ответа
-    title="Банковский чат-бот",  # Заголовок страницы
-    description="Введите вопрос клиента — получите несколько ответов от банка."  # Описание
+    fn=generate_response,
+    inputs=gr.Textbox(lines=2, placeholder="Введите банковский вопрос..."),
+    outputs="text",
+    title="Русский банковский чат-бот",
+    description="Задайте вопрос, например: 'Какие комиссии за переводы?' или 'Как активировать карту?'"
 )
 
-# Запуск интерфейса
 iface.launch()
 
 
-