File size: 820 Bytes
bed8885
 
 
9d1d724
bed8885
 
 
 
 
43dd320
7d94a8b
 
9d1d724
bed8885
49ba46b
 
 
bed8885
9d1d724
 
49ba46b
 
 
 
9d1d724
bed8885
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
'''
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(' ','') 

iface = gr.Interface(
  fn=poem, 
  inputs=[
    gr.Textbox(lines=1, default="春眠不觉晓,",label="请输入一句诗词:"),
    gr.Number(50,label='字数:',default=50),
  ],
  outputs="text")
iface.launch()