TransformersでGPTの文章生成する方法いろいろ

Transformersを使うと、GPTの事前学習モデルを使って簡単に文章生成ができます。モデル自体は同じでも色々なメソッドが用意されていて、用途に応じて適切なインターフェースを選ぶことでより便利に使えます。

## 環境
- Google Colaboratory
- transformers: 4.25.1

`!pip install transformers sentencepiece`等でtransformersをインストールしておきます。

### 事前準備
`MODEL_NAME`定数に好きなモデルを指定してください。ここでは[rinna/japanese\-gpt2\-xsmall](https://huggingface.co/rinna/japanese-gpt2-xsmall)を使います。

```python
MODEL_NAME = 'rinna/japanese-gpt2-xsmall'
```

## pipeline
`pipeline()`を使うのが、最も簡単な方法だと思います。
`pipeline()`の第一引数にタスクを指定することで、タスクを実行する簡単なパイプラインを生成できます。。ここではテキスト生成なので`text-generation`を用いますが、他に使えるタスクは[公式ドキュメント](https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.pipeline.task)に説明があります。

```python
from transformers import pipeline

text_pipe = pipeline('text-generation', model=MODEL_NAME)
output = text_pipe("昔々あるところに")

output[0]['generated_text']
```

>昔々あるところに、たくさん人が来て、たくさんの人がいて、みんなで楽しめる、それがなんだろう...と思ったことがあったんですw 僕には「この人じゃなかったらやらなかったんだねー」と思いました。 ・さん:

※この出力は毎回変わります

## `generate()`メソッド
おそらく最も一般的なのは`generate()`メソッドを使った方法ではないでしょうか。[rinnaモデルのREADME](https://huggingface.co/rinna/japanese-gpt-1b/blob/main/README.md)で紹介されている方法だったりします。

```python
import torch
from transformers import T5Tokenizer, AutoModelForCausalLM

tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

with torch.no_grad():
    input_ids = tokenizer.encode("昔々あるところに", add_special_tokens=False, return_tensors="pt")
    output_ids = model.generate(input_ids.to(model.device), max_length=20, pad_token_id=tokenizer.pad_token_id)

tokenizer.decode(output_ids.tolist()[0])
```

>昔々あるところに、あるお店があります。 店内は、カウンター席とテーブル


`generate()`に引数として`top_k`や`num_beams`などを付与すると、beam searchやsamplingなどのdecodingアルゴリズムを制御できます。decodingアルゴリズムについて詳細は下記公式ブログを参照してください。

https://huggingface.co/blog/how-to-generate

## greedy searchを自前で実装する
decodingの中でも最も単純なgreedy search（貪欲法）つまり、もっとも確率の高いトークンを選び続ける手法を、実際に実装してみます。
なお、これは[機械学習エンジニアのためのTransformers](https://www.oreilly.co.jp/books/9784873119953/)にて紹介されていたソースコードを利用しています。

```python
n_steps = 10

input_ids = tokenizer.encode("昔々あるところに", return_tensors="pt")

with torch.no_grad():
    for _ in range(n_steps):
        output = model(input_ids)

        next_token_logits = output.logits[0, -1, :]
        next_token_probs = torch.softmax(next_token_logits, dim=-1)
        sorted_ids = torch.argsort(next_token_probs, descending=True, dim=-1)

        input_ids = torch.cat([input_ids, sorted_ids[None, 0, None]], dim=-1)

tokenizer.decode(input_ids[0])
```

```:出力
昔々あるところに
<unk> <unk> <unk> <unk> <unk>
```

unknown token `<unk>` ばかりになってしまいました。実装を間違えてしまったのでしょうか。実際に、モデルの出力がどうなっているのかを確認してみましょう。トークンを確率が高い順に並び替えた結果である`sorted_ids`の中身を出力してみます。


```diff_python
n_steps = 10

input_ids = tokenizer.encode("昔々あるところに", return_tensors="pt")

with torch.no_grad():
    for _ in range(n_steps):
        output = model(input_ids)

        next_token_logits = output.logits[0, -1, :]
        next_token_probs = torch.softmax(next_token_logits, dim=-1)
        sorted_ids = torch.argsort(next_token_probs, descending=True, dim=-1)

        input_ids = torch.cat([input_ids, sorted_ids[None, 0, None]], dim=-1)

+       print(",".join(tokenizer.decode(sorted_ids[i]) for i in range(5)))

```

```:出力
<unk>,あ,昔,こんな,は
,C,の,o,oo
<unk>,が,の,という,で
,C,が,の,という
<unk>,が,の,という,C
,C,が,の,という
<unk>,が,という,の,C
,C,が,の,(
<unk>,が,という,の,C
,C,が,(,の
```

第1位に`<unk>`が来ているのですが、第2位以下は普通の単語が並んでいる事がわかります。

普通こんな事はしませんが、第1位ではなく第2位のトークンを取り続けるように改修すると、文章を出力するようになりました。

```diff_python
n_steps = 10

input_ids = tokenizer.encode("昔々あるところに", return_tensors="pt")

with torch.no_grad():
    for _ in range(n_steps):
        output = model(input_ids)

        next_token_logits = output.logits[0, -1, :]
        next_token_probs = torch.softmax(next_token_logits, dim=-1)
        sorted_ids = torch.argsort(next_token_probs, descending=True, dim=-1)

-       input_ids = torch.cat([input_ids, sorted_ids[None, 0, None]], dim=-1)
+       input_ids = torch.cat([input_ids, sorted_ids[None, 1, None]], dim=-1)

tokenizer.decode(input_ids[0])
```

```:出力
昔々あるところに
 あそこは、この辺で ”<unk>
```

このように、自前で実装することで中身の挙動を調べたりカスタムが自由に行なえます。

### `greedy_search()`メソッド
greedy searchを自前で実装しましたが、Transformers内部にもメソッドが用意されています。
このメソッドを使うことで処理自体はとても単純になるのですが、`LogitsProcessor`というものを用意して渡す必要があります。説明は置いておいて、とりあえず動かしてみます。

```python
from transformers.generation import LogitsProcessorList

input_ids = tokenizer.encode("昔々あるところに", return_tensors="pt")

logits_processor = model._get_logits_processor(
    repetition_penalty=None,
    no_repeat_ngram_size=None,
    encoder_no_repeat_ngram_size=None,
    input_ids_seq_length=None,
    encoder_input_ids=None,
    bad_words_ids=None,
    min_length=None,
    max_length=20,
    eos_token_id=tokenizer.eos_token_id,
    forced_bos_token_id=tokenizer.bos_token_id,
    forced_eos_token_id=tokenizer.eos_token_id,
    prefix_allowed_tokens_fn=None,
    num_beams=None,
    num_beam_groups=None,
    diversity_penalty=None,
    remove_invalid_values=None,
    exponential_decay_length_penalty=None,
    logits_processor=LogitsProcessorList(),
    renormalize_logits=None,
)

outputs = model.greedy_search(
    input_ids=input_ids,
    logits_processor=logits_processor,
    pad_token_id=tokenizer.pad_token_id,
)

tokenizer.decode(outputs[0])
```

```:出力
昔々あるところに
<unk> <unk> <unk> <unk> <unk> <unk> <unk>

```

greedy searchを実装した時と同様、`<unk>`が並んでしまいました。

`LogitsProcessor`には`bad_words_ids`という設定があり、ここに`<unk>`を出力しないよう指定することができます。

```diff_python
logits_processor = model._get_logits_processor(
    repetition_penalty=None,
    no_repeat_ngram_size=None,
    encoder_no_repeat_ngram_size=None,
    input_ids_seq_length=None,
    encoder_input_ids=None,
-   bad_words_ids=None,
+   bad_words_ids=[[tokenizer.unk_token_id]],
    min_length=None,
    max_length=20,
    eos_token_id=tokenizer.eos_token_id,
    forced_bos_token_id=tokenizer.bos_token_id,
    forced_eos_token_id=tokenizer.eos_token_id,
    prefix_allowed_tokens_fn=None,
    num_beams=None,
    num_beam_groups=None,
    diversity_penalty=None,
    remove_invalid_values=None,
    exponential_decay_length_penalty=None,
    logits_processor=LogitsProcessorList(),
    renormalize_logits=None,
)
```

```:出力
昔々あるところに
 ああ、あの頃は ああ、あの頃は

```

このように、`LogitsProcessor`を使うことで必要な出力が得やすいようにdecodingできます。

## 自前の実装をカスタムする
### `LogitsProcessor`を利用する
先ほど自前で実装したgreedy searchに`LogitsProcessor`を組み込むには、下記のようにモデルの出力に挟みます。

```python
input_ids = tokenizer.encode("昔々あるところに", return_tensors="pt")
n_steps = 10

with torch.no_grad():
    for _ in range(n_steps):
        output = model(input_ids)
        next_token_logits = output.logits[:, -1, :]
        next_token_scores = logits_processor(input_ids, next_token_logits)

        probs = torch.softmax(next_token_scores, dim=-1)

        next_tokens = torch.argmax(probs, dim=-1)[:, None]

        input_ids = torch.cat([input_ids, next_tokens], dim=-1)

tokenizer.decode(input_ids[0])
```

```:出力
昔々あるところに
ああ、あの頃は ああ、
```

ちゃんと、先ほどと同じ出力が得られました。
このようにして、自前の実装とTransformersの便利なメソッドを組み合わせることで、必要な生成処理を比較的簡単に実装できるようになります。

### samplingを実装する
これまではgreedy searchということで、`torch.argmax()`を用いて確率が最大のトークンを持ってきていました。ここを変更し`torch.multinomial`で確率分布に基づいてトークンを選ぶようにすると、samplingになります。

```diff_python
input_ids = tokenizer.encode("昔々あるところに", return_tensors="pt")
n_steps = 10

with torch.no_grad():
    for _ in range(n_steps):
        output = model(input_ids)
        next_token_logits = output.logits[:, -1, :]
        next_token_scores = logits_processor(input_ids, next_token_logits)

        probs = torch.softmax(next_token_scores, dim=-1)

-       next_tokens = torch.argmax(probs, dim=-1)[:, None]
+       next_tokens = torch.multinomial(probs, num_samples=1)[:, None, 0]

        input_ids = torch.cat([input_ids, next_tokens], dim=-1)

tokenizer.decode(input_ids[0])
```

```:出力 ※この出力は毎回変わります
昔々あるところに
いろんな種類校正された校正物
```

同様に、この処理を変更していけばbeam searchやtop-k samplingなども実装できます。

## まとめ
TransfomrmersでGPTの文章生成する方法を複数紹介しました。手軽な方法を使えばすぐに文章生成を試すことができますし、自前で実装することで仕組みの理解を深めたり、調査やカスタマイズもできるようになります。Transfomrmersには色々な機能が用意されていて、目的に応じて適切な機能を使いこなす事が大切だと思います。

## 参考文献
- [O'Reilly Japan \- 機械学習エンジニアのためのTransformers](https://www.oreilly.co.jp/books/9784873119953/)
- [huggingface/transformersのgeneration_utils](https://github.com/huggingface/transformers/blob/cc3d0e1b017dbb8dcbba1eb01be77aef7bacee1a/src/transformers/generation/utils.py)をひたすら読んだ