import gradio as gr
from transformers import pipeline

generator = pipeline('text-generation', model='huggingtweets/wwm_shakespeare')


def greet(input_words):
    return generator(input_words)[0]['generated_text']


iface = gr.Interface(fn=greet, inputs="text", outputs="text", layout='vertical',
                     title='Shakespeare Love Notes Generator',
                     description='Input some words, and generate some love notes in the writing style of William Shakespeare. Check the Linked Models for more detail.',
                     examples=['Chris', 'Love You', 'Special Day', 'Speak Love', 'Green Mountain', "Valentine's Day"], article='**Keep clicking Submit for more results with the same input**')
iface.launch()