import gradio as gr
from transformers import pipeline

# Set up the generatove model transformer pipeline
generator = pipeline("text-generation", model="matthh/gpt2-rlhf-joyous-poetry")


def generate(prompt):
    result = generator(prompt, max_length=50, num_return_sequences=1)
    print(result)
    return result[0]["generated_text"]


def reset_all(new_line):
    return None


demo = gr.Blocks()

with demo:
    with gr.Row():
        with gr.Column():
            input = gr.Textbox(
                lines=1,
                value="Music to hear, why hear'st thou music sadly?\nSweets with sweets war not, joy delights in joy.",
            )
            with gr.Row():
                generate_button = gr.Button("Generate", variant="primary")
                clear_all_button = gr.Button("Reset")
        with gr.Column():
            output = gr.HTML()

    generate_button.click(fn=generate, inputs=input, outputs=output)
    clear_all_button.click(fn=reset_all, inputs=input, outputs=output)

if __name__ == "__main__":
    demo.launch()