import gradio as gr
from transformers import pipeline

# Загрузка модели для генерации текста
generator = pipeline("text-generation", model="gpt2")

# CSS стили
css = """
#generate {
    width: 100%;
    background: #e253dd !important;
    border: none;
    border-radius: 50px;
    outline: none !important;
    color: white;
}
#generate:hover {
    background: #de6bda !important;
    outline: none !important;
    color: #fff;
}
#image_output {
    display: flex;
    justify-content: center;
}
footer {visibility: hidden !important;}
#image_output {
    height: 100% !important;
}
"""

# Функция для генерации текста по prompt
def generate_text(prompt):
    output = generator(prompt, max_length=100)[0]['generated_text']
    return output

# Создание интерфейса Gradio
demo = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    layout="vertical",
    title="Text Generation Demo",
    css=css
)

# Запуск интерфейса
if __name__ == "__main__":
    demo.launch()