skytnt
/

gpt2-japanese-lyric-medium

Text Generation

text-generation-inference

Inference Endpoints

Model card Files Files and versions Community

skytnt commited on Jul 9, 2022

Commit

3d1a488

•

1 Parent(s): 4e1316c

Update README.md

Files changed (1) hide show

README.md +5 -1

README.md CHANGED Viewed

@@ -24,9 +24,13 @@ The model is used to generate Japanese lyrics.
 import torch
 from transformers import T5Tokenizer, GPT2LMHeadModel
 tokenizer = T5Tokenizer.from_pretrained("skytnt/gpt2-japanese-lyric-medium")
 model = GPT2LMHeadModel.from_pretrained("skytnt/gpt2-japanese-lyric-medium")
 def gen_lyric(title: str, prompt_text: str):
     if len(title)!= 0 or len(prompt_text)!= 0:

 import torch
 from transformers import T5Tokenizer, GPT2LMHeadModel
+device = torch.device("cpu")
+if torch.cuda.is_available():
+    device = torch.device("cuda")
 tokenizer = T5Tokenizer.from_pretrained("skytnt/gpt2-japanese-lyric-medium")
 model = GPT2LMHeadModel.from_pretrained("skytnt/gpt2-japanese-lyric-medium")
+model = model.to(device)
 def gen_lyric(title: str, prompt_text: str):
     if len(title)!= 0 or len(prompt_text)!= 0: