import gradio as gr
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Загрузка модели ruT5 и токенизатора
model_name = "cointegrated/rut5-base-multi-sentence-task"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def generate_text(input_text):
    # Применение токенизатора к входному тексту
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    
    # Генерация текста на основе входных данных
    output = model.generate(input_ids)
    
    # Декодирование сгенерированного текста
    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
    
    return decoded_output

# Создание интерфейса Gradio
input_text = gr.inputs.Textbox(lines=5, label='Введите текст для генерации')
output_text = gr.outputs.Textbox(label='Сгенерированный текст')
interface = gr.Interface(fn=generate_text, inputs=input_text, outputs=output_text)

# Запуск интерфейса
interface.launch()