Spaces:
Running
Running
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() | |