| ''' | |
| 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) | |
| title = "写一首诗" | |
| article = "<div style='text-align: center;'>使用的模型:<a href='https://huggingface.co/uer/gpt2-chinese-poem' target='_blank'>https://huggingface.co/uer/gpt2-chinese-poem</a></div>" | |
| 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("今晚提测又加班，BUG一堆改不完。", lines=1, label="请输入第一句:"), | |
| gr.Number(66, label='字数:'), | |
| ], | |
| outputs=gr.Textbox(lines=6), | |
| title=title, | |
| article=article) | |
| iface.launch() |