import gradio as gr

from transformers import BertTokenizer, GPT2LMHeadModel,TextGenerationPipeline
tokenizer = BertTokenizer.from_pretrained("uer/gpt2-chinese-poem")
model = GPT2LMHeadModel.from_pretrained("uer/gpt2-chinese-poem")
text_generator = TextGenerationPipeline(model, tokenizer) 

def predict(prompt):
    completion = text_generator(prompt, max_length=50, do_sample=True)
    return completion
    
gr.Interface(fn=predict, inputs="text", outputs="text").launch()
#def greet(name):
#    return "Hello " + name + "!!"

#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
#iface.launch()