chat / app.py
SimrusDenuvo's picture
Update app.py
f1cded5 verified
raw
history blame
2.09 kB
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
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=generate_response,
inputs=gr.Textbox(lines=2, placeholder="Введите банковский вопрос..."),
outputs="text",
title="Русский банковский чат-бот",
description="Задайте вопрос, например: 'Какие комиссии за переводы?' или 'Как активировать карту?'"
)
iface.launch()