chat / app.py
SimrusDenuvo's picture
Update app.py
90cd180 verified
raw
history blame
1.57 kB
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# Загружаем модель и токенизатор
model_name = "ai-forever/rugpt3small_based_on_gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Функция генерации ответа
def generate_response(prompt):
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=150,
do_sample=True,
temperature=0.7,
top_k=50,
top_p=0.95,
eos_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Удаляем промпт из ответа, если повторяется
if response.startswith(prompt):
response = response[len(prompt):].strip()
return response.strip()
# Интерфейс Gradio
demo = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(lines=4, label="Введите вопрос по клиентским обещаниям банка"),
outputs=gr.Textbox(label="Ответ модели"),
title="Анализ клиентских обещаний — RuGPT-3",
description="Используется модель ai-forever/rugpt3small_based_on_gpt2 на основе данных ZhenDOS/alpha_bank_data."
)
# Запуск
if name == "main":
demo.launch()