from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import gradio as g

# Загружаем модель rut5-base с медленным токенизатором
model_name = "cointegrated/rut5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Функция генерации ответа
def generate_response(prompt):
    input_ids = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=200,
            num_beams=5,
            do_sample=True,
            top_p=0.9,
            temperature=0.8,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return response

# Gradio-интерфейс
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=4, label="Введите ваш запрос"),
    outputs=gr.Textbox(label="Ответ модели"),
    title="Интерфейс ChatGPT",
    description="Пример взаимодействия с API OpenAI через Hugging Face Space"
)

iface.launch()