File size: 1,491 Bytes
5aeff64 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import gradio as gr
import torch
from transformers import T5Tokenizer, AutoModelForCausalLM
from utils import translate_from_jp_to_en
tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt-1b")
model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-1b")
def generate(text, max_length=128):
token_ids = tokenizer.encode(
text, add_special_tokens=False, return_tensors="pt")
with torch.no_grad():
output_ids = model.generate(
token_ids,
max_length=max_length,
do_sample=True,
top_k=500,
top_p=0.95,
pad_token_id=tokenizer.pad_token_id,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id,
bad_word_ids=[[tokenizer.unk_token_id]]
)
output = tokenizer.decode(output_ids.tolist()[0])
return output, translate_from_jp_to_en(output)
title = "JP GPT Demo"
description = "Demo for generating text in Japanase using a GPT model"
examples = [['日本のeスポーツ障害者がステレオタイプを撃ち落とす', 128]]
gr.Interface(fn=generate, inputs=[gr.inputs.Textbox(lines=4, label="Prompt"),
gr.inputs.Slider(minimum=8, maximum=1024, step=8, default=64, label="Numbers of tokens")],
outputs=["text", "text"],
title=title, description=description,
# article= article,
examples=examples).launch(enable_queue=True)
|