Transformaers Pipeline によるテキスト生成

Hugging Face の Transformers には推論を簡単に行うための Pipeline という仕組みがあります。 PyTorch や TensorFlow のような面倒なコードを書かずに、わずかなコーディング量で推論することができるので、ちょっとしたことを試すのにはとでも便利です。

Pipeline 自体は画像認識や音声認識など様々なタスクに使用できるようですが、今回は自然言語処理のタスク、特にテキスト生成を試してみます。

特定のモデルで Pipeline を使用する流れはこんな感じです。

1. 学習済みモデルをロード
2. タスクやモデルなどを指定して Pipeline を構築
3. 推論の実行（パイプラインに入力を与えて出力を得る）

ちなみに、最も単純な場合は、明示的に学習済みモデルをロードせずにタスクを指定するだけでも良さそうですが、使用可能な言語が限られるなどの制限がありそうです。 今回は、最近公開された日本語の大規模言語モデルを使用して、日本語テキスト生成を試したかったので、上記のような流れで実装してみます。

なお、今回のコードはこちらです。

github.com

学習済みモデルをロード

学習済みモデルとして、今回は、サイバーエージェント社が公開している日本語データセットで学習済みの OpenCALM-1B を使用してみます。

from transformers import AutoModelForCausalLM, AutoTokenizer

pretrained_model_name = 'cyberagent/open-calm-1b'

model = AutoModelForCausalLM.from_pretrained(pretrained_model_name, torch_dtype=torch.float16).to(device)
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name)

パイプラインを構築

続いてパイプラインを構築する。 テキスト生成の場合は transformers.pipeline の引数 task に text-generation を指定し、model と tokenizer に、先ほどロードした学習済みモデルとトークナイザーを指定する。 kwargs の値は OpenCALM-1B の Usage を参考にしました。

from transformers import pipeline

task = 'text-generation'

kwargs = {
    'max_new_tokens': 64,
    'do_sample': True,
    'temperature': 0.7,
    'top_p': 0.9,
    'repetition_penalty': 1.05,
    'pad_token_id': tokenizer.pad_token_id
}

generator = pipeline( task=task, model=model, tokenizer=tokenizer, device=device_id, torch_dtype=torch.float16, **kwargs )

推論の実行

引数にテキストを与えると、それに続くテキストが生成されます。 たったのこれだけ。なんて簡単なんだ！

generator('アジャイルソフトウェア開発宣言の内容は')

ちなみに出力結果はこんな感じでした。

[{'generated_text': 'アジャイルソフトウェア開発宣言の内容は、ソフトウェアの開発・運用に関する方針を簡潔にまとめたものです。\nこの文書には、「開発」や「生産」「販売」、「保守サービスなどの各業務プロセスにおいて『顧客満足』を追求していくこと」、そしてそれらを実現するための考え方として、『SMILE(Smile:笑顔で)』と書かれています。『信頼』『誠実さ'}]

使用する学習済みモデルやモデルのパラメータ数によって精度も変わると思います。 様々なモデルで試してみたいですね。