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()