| ''' | |
| 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( | |
| gr.Markdown( | |
| """ | |
| # 写一首诗 | |
| 模型：https://huggingface.co./uer/gpt2-chinese-poem | |
| """), | |
| fn=poem, | |
| inputs=[ | |
| gr.Textbox("今晚提测又加班，BUG一堆改不完。", lines=1, label="请输入一句诗词:"), | |
| gr.Number(66, label='字数:'), | |
| ], | |
| outputs=gr.Textbox(lines=6)) | |
| iface.launch() |