Update README.md

README.md

@@ -31,10 +31,10 @@ model = GPT2LMHeadModel.from_pretrained("skytnt/gpt2-japanese-lyric-small")
 
 
 def gen_lyric(prompt_text: str):
-    prompt_text = prompt_text.replace("\n", "\\n ")
+    prompt_text = "<s>" + prompt_text
     prompt_tokens = tokenizer.tokenize(prompt_text)
     prompt_token_ids = tokenizer.convert_tokens_to_ids(prompt_tokens)
-    prompt_tensor = torch.LongTensor(prompt_token_ids)
+    prompt_tensor = torch.LongTensor(prompt_token_ids).to(device)
     prompt_tensor = prompt_tensor.view(1, -1)
     # model forward
     output_sequences = model.generate(
@@ -55,8 +55,7 @@ def gen_lyric(prompt_text: str):
     generated_sequence = output_sequences.tolist()[0]
     generated_tokens = tokenizer.convert_ids_to_tokens(generated_sequence)
     generated_text = tokenizer.convert_tokens_to_string(generated_tokens)
-    generated_text = "\n".join([s.strip() for s in generated_text.split('\\n')]).replace(' ', '\u3000').replace(
-        '</s>', '\n\n---end---')
+    generated_text = "\n".join([s.strip() for s in generated_text.split('\\n')]).replace(' ', '\u3000').replace('<s>', '').replace('</s>', '\n\n---end---')
     return generated_text