| ''' | |
| Author: zuojianghua | |
| Date: 2022-07-06 10:40:08 | |
| LastEditTime: 2022-07-07 17:03:34 | |
| LastEditors: zuojianghua | |
| Description: | |
| FilePath: /zuo/app.py | |
| ''' | |
| import gradio as gr | |
| from transformers import BertTokenizer, TFGPT2LMHeadModel,TextGenerationPipeline | |
| tokenizer = BertTokenizer.from_pretrained(".") | |
| model = TFGPT2LMHeadModel.from_pretrained(".") | |
| text_generator = TextGenerationPipeline(model, tokenizer) | |
| def poem(cls, max_length): | |
| txt = text_generator('[CLS]'+cls, max_length=int(max_length), do_sample=True) | |
| return txt[0]['generated_text'].replace('[CLS]','').replace('[SEP]','').replace(' ','').replace('。','。\n') | |
| iface = gr.Interface( | |
| fn=poem, | |
| inputs=[ | |
| gr.Textbox(lines=1, default="春眠不觉晓，",label="请输入一句诗词:"), | |
| gr.Number(50, label='字数:',default=50), | |
| ], | |
| outputs=gr.Textbox(lines=4)) | |
| iface.launch() |