from transformers import pipeline
import gradio as gr

pipe = pipeline('text-generation', model='daspartho/prompt-extend')

def extend_prompt(prompt):
    return pipe(prompt+',', num_return_sequences=1)[0]["generated_text"]

iface = gr.Interface(
    fn=extend_prompt, 
    inputs=gr.Text(label="Type the prompt here"), 
    outputs=gr.TextArea(label='Extended prompt'),
    title="PROMPT EXTENDER",
    examples=[
        ["Write a creative story about a robot learning to paint."],
        ["Generate a poem about the beauty of nature."],
        ["Create a dialogue between a detective and a suspect."],
        ["Imagine a futuristic city where technology and art blend seamlessly."]
    ]
)

iface.launch()