zuo / app.py
zuozuo's picture
Update app.py
43dd320
raw
history blame
670 Bytes
'''
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("./tokenizer_config.json")
model = TFGPT2LMHeadModel.from_pretrained("./config.json")
text_generator = TextGenerationPipeline(model, tokenizer)
def poem(cls):
txt = text_generator('[CLS]'+cls, max_length=100, do_sample=True)
return txt
iface = gr.Interface(
fn=poem,
inputs=gr.Textbox(lines=1, placeholder="春眠不觉晓,"),
outputs="text")
iface.launch()