Huggingface Transformersを使って10行で桃太郎の続きを生成する

5654 ワード
Python transformers 自然言語処理 Python テキストリンク

動作環境

  • Google Colaboratory
  • Huggingface Transformers 2.5.1
  • Pytorch 1.5.0

実装内容

  • 「桃太郎」でおばあさんが桃を割ったところまでの文章を入力し、その続きを日本語版BERTに予測させる。

ソースコード

generate_ja_sentence.ipynb
!pip install torch  torchvision -f https://download.pytorch.org/whl/torch_stable.html
!pip install transformers
!pip install fugashi
!pip install ipadic
import torch
from transformers import AutoModelWithLMHead, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking") 
model = AutoModelWithLMHead.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")

prompt = "昔々、あるところにおじいさんとおばあさんが住んでいました。おじいさんは山へ芝刈りに、おばあさんは川で洗濯にいきました。おばあさんが川で洗濯をしていると、大きな桃が流れてきました。おばあさんは大きな桃を拾いあげて家に持ち帰りました。そしておじいさんとおばあさんが桃を食べようと桃を切ってみると、"

inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
prompt_length = len(tokenizer.decode(inputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True))
outputs = model.generate(inputs, max_length=150, do_sample=True, top_k=60, top_p=0.98)
generated = tokenizer.decode(outputs[0])[prompt_length:].replace(" ", "")

print(prompt + generated)

生成した文章はこのようになりました。

昔々、あるところにおじいさんとおばあさんが住んでいました。おじいさんは山へ芝刈りに、おばあさんは川で洗濯にいきました。おばあさんが川で洗濯をしていると、大きな桃が流れてきました。おばあさんは大きな桃を拾いあげて家に持ち帰りました。そしておじいさんとおばあさんが桃を食べようと桃を切ってみると、偶然みてまたこの桃が降りてくる。すると、桃が降りてくる。するとまた桃が割れる。すると、桃が割れてしまうので、もう一度少しずつ割れ続ける。すると、おばあさんが言う

めちゃくちゃホラー展開になってしまいました。

まとめ

上手く生成できたとは言い難いですが、実質数行で文章生成が出来てしまいました。
まだまだ勉強不足なので、この数行の中身を理解するまでは時間がかかりそうです
ここまで読んでいただきありがとうございました。

参考リンク

Huggingface Transformers 入門 (1)