| 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() | |