|
from spaces import GPU |
|
import torch |
|
from transformers import pipeline, AutoTokenizer |
|
import gradio as gr |
|
|
|
|
|
model_name = "rinna/japanese-gpt2-medium" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name, legacy=False) |
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
generator = pipeline('text-generation', model=model_name, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1) |
|
|
|
@GPU(duration=120) |
|
|
|
def generate_text(prompt, max_length): |
|
result = generator(prompt, max_length=max_length, num_return_sequences=1) |
|
return result[0]['generated_text'] |
|
|
|
iface = gr.Interface( |
|
fn=generate_text, |
|
inputs=[ |
|
gr.Textbox(label="プロンプト", placeholder="ここに日本語のプロンプトを入力してください"), |
|
gr.Slider(minimum=10, maximum=20000, value=50, step=1, label="最大長") |
|
], |
|
outputs=gr.Textbox(label="生成されたテキスト") |
|
) |
|
|
|
iface.launch(share=True) |