メモ帳

python, juliaで機械学習をやっていく

(Part. 2) tensorflow 2.0 betaでtransformerをつかって言語生成chatbotをつくりたい

以下の記事のPart. 2です。 Part. 1ではtransformerの説明とモデルの実装をしました。 この記事ではlossやmetricsを定義し実際に学習を行います。 また、日本語データを用いるために分かち書きも実装しています。

tksmml.hatenablog.com

import

import tensorflow as tf

loss

lossはcategorical crossentropyです。padding部分(y_true=0)は無視します。

def loss_function(y_true, y_pred):
    y_true = tf.reshape(y_true, shape=(-1, MAX_LENGTH - 1))

    loss = tf.keras.losses.SparseCategoricalCrossentropy(
      from_logits=True, reduction='none')(y_true, y_pred)

    # y_true!=0 -> 1.0, y_true==0 -> 0.0
    # に変換し、lossに乗算し、y_true=0部分のlossをすべて0.0にする
    mask = tf.cast(tf.not_equal(y_true, 0), tf.float32)
    loss = tf.multiply(loss, mask)

    return tf.reduce_mean(loss)

optimizer

Adamで以下のスケジューラーを用います。

 l_{rate} = {d-model}^{-0.5} \cdot min(step-num^{-0.5}, step-num \cdot warmup-steps^{-1.5})

論文でも言及されていますが、学習データが少ない場合は有効に働かないようです。その場合はデフォルトのAdamを使ってしまっていいと思います。

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):

    def __init__(self, d_model, warmup_steps=4000):
        super(CustomSchedule, self).__init__()

        self.d_model = tf.cast(d_model, tf.float32)
        self.warmup_steps = warmup_steps

    def __call__(self, step):
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warmup_steps**-1.5)

        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

metrics

各単語の一致率を正解率とします。

def accuracy(y_true, y_pred):
    # ensure labels have shape (batch_size, MAX_LENGTH - 1)
    y_true = tf.reshape(y_true, shape=(-1, MAX_LENGTH - 1))
    accuracy = tf.metrics.SparseCategoricalAccuracy()(y_true, y_pred)
    return accuracy

compile

以上の内容でモデルを定義し、コンパイルします。 transformer modelはPart. 1で定義したものです。

tf.keras.backend.clear_session()

NUM_LAYERS = 2
EMB_DIM = 256
NUM_HEADS = 8
UNITS = 512
DROPOUT = 0.1

model = transformer(
    vocab_size=VOCAB_SIZE,
    num_layers=NUM_LAYERS,
    units=UNITS,
    emb_dim=EMB_DIM,
    head=NUM_HEADS,
    dropout=DROPOUT)


learning_rate = CustomSchedule(EMB_DIM)

optimizer = tf.keras.optimizers.Adam(
    learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)

model.compile(optimizer=optimizer, loss=loss_function, metrics=[accuracy])

学習データ

使用するデータは自前で用意したもなので、公開できませんが日本語の1対1の対話データです。

分かち書き

日本語だとそのままでは学習に使えないので、分かち書きを行い、tokenに分割します。 分かち書きMeCabを使います。他にも、Janomeであれば導入が簡単なので、そちらもオススメです。 collaboratoryへのインストールが必要です。以下を実行して下さい。

!apt install aptitude
!aptitude install mecab libmecab-dev mecab-ipadic-utf8 git make curl xz-utils file -y
!pip install mecab-python3==0.7

インストールが成功しているか確認します。

import MeCab
import re


def preprocess_sentence(sentence):
    tagger = MeCab.Tagger("-Owakati")
    parser = tagger.parse(sentence)
    return re.sub('\n', '', parser)


preproces_sentence("隣の客はよく柿食う客だ")
# >>> '隣 の 客 は よく 柿 食う 客 だ'

あらかじめデータセットの文書を分かち書きしておきます。 (predict時にも上記関数を噛まして分かち書きしたものを入力することを忘れないようにします。)

学習データセットは以下のようなテキストファイル想定です。同じ行数のuser.txtとsystem.txtが1対1の会話セットになっています。

user.txt

こんにちは、はじめまして
暑いですね

system.txt

こんにちは、こちらこそはじめまして
クーラーの効いている部屋に行きたいですね

上記データセット分かち書きした文章をリスト形式で保持します。

with open('user.txt', 'r') as f:
    user_data = f.read().split('\n')
    user_data = list(map(preprocess_sentence, user_data))

with open('system.txt', 'r') as f:
    sys_data = f.read().split('\n')
    sys_data = list(map(preprocess_sentence, sys_data))

print(user_data)
>>> ['こんにちは 、 はじめまして', '暑い です ね']
print(sys_data)
>>> ['こんにちは 、 こちら こそ はじめまして', 'クーラー の 効い て いる 部屋 に 行き たい です ね']

mecabに関連するコードは以上です。janomeなどの他の形態素解析ライブラリを使った場合でも形式をそろえると以下は本記事と同様なコードでも動くと思います。

tensorflow dataset

MeCabが使えるようになったので、次はtensorflowで前処理のコードを書いていきます。 tensorflow datasetというものを使えば高パフォーマンスの入力パイプラインが定義できます。 まず、以下を実行して準備します。

!pip install tfds-nightly


import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow.keras.preprocessing.sequence import pad_sequences

まず、学習データから単語IDの辞書を作成します。 また、transformerでは文字列の最初と最後に固有のIDを挿入します。そのためにあらかじめ固有のIDを割り振ります。 そして、embedding layer用に学習データの単語数を取得します。

# user発話とsytem発話両方を使ってtokenizerの作成
tokenizer = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    user_data + sys_data, target_vocab_size=2**13)

# 文字列の最初と最後に固有のIDを定義
START_TOKEN, END_TOKEN = [tokenizer.vocab_size], [tokenizer.vocab_size + 1]

# ボキャブラリサイズに最初と最後を表すIDの数を足す
VOCAB_SIZE = tokenizer.vocab_size + 2

次に、前処理を行う関数を定義します。主に以下の機能をもちます。

  • 単語をIDに変換する
  • 最大文字列を指定して0埋めする
  • 文字列の最初と最後にstart ID, end IDを割り振る
  • 最大文字列をこえる学習データを省く
# Maximum sentence length
MAX_LENGTH = 40


def tokenize_and_filter(inputs, outputs, tokenizer, MAX_LENGTH, START_TOKEN, END_TOKEN):
    tokenized_inputs, tokenized_outputs = [], []
    to_inp_a, to_out_a = tokenized_inputs.append, tokenized_outputs.append
  
    for (sentence1, sentence2) in zip(inputs, outputs):
        # tokenize sentence
        sentence1 = START_TOKEN + tokenizer.encode(sentence1) + END_TOKEN
        sentence2 = START_TOKEN + tokenizer.encode(sentence2) + END_TOKEN
        # max sentence lengthをこえていないデータだけ使用
        if len(sentence1) <= MAX_LENGTH and len(sentence2) <= MAX_LENGTH:
            to_inp_a(sentence1)
            to_out_a(sentence2)
  
    # token化されたsentenceを0埋めする
    tokenized_inputs = pad_sequences(tokenized_inputs, maxlen=MAX_LENGTH, padding='post')
    tokenized_outputs = pad_sequences(tokenized_outputs, maxlen=MAX_LENGTH, padding='post')
  
    return tokenized_inputs, tokenized_outputs


user, sys = tokenize_and_filter(user_data, sys_data, tokenizer, MAX_LENGTH, START_TOKEN, END_TOKEN)

tensorflow用に名前に注意して、tensorflow dataの形式にする。

BATCH_SIZE = 64
BUFFER_SIZE = 20000

# decoderの入力には前のtargetを使う
# START_TOKENをtargetから削除
dataset = tf.data.Dataset.from_tensor_slices((
    {
        'inputs': user,
        'decoder_inputs': sys[:, :-1]
    },
    {
        'outputs': sys[:, 1:]
    },
))

dataset = dataset.cache()
dataset = dataset.shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

前処理用のコードは以上です。

prediction

学習経過を見るために、エポック毎に言語生成することにします。そのためにprediction関数を定義しておきます。

def evaluate(sentence, tokenizer, model, START_TOKEN, END_TOKEN, MAX_LENGTH):
    sentence = preprocess_sentence(sentence)
    sentence = tf.expand_dims(
      START_TOKEN + tokenizer.encode(sentence) + END_TOKEN, axis=0)

    output = tf.expand_dims(START_TOKEN, 0)

    # modelがEND_TOKENを予測するまで繰り返しpredictさせる
    for i in range(MAX_LENGTH):
        predictions = model(inputs=[sentence, output], training=False)

        # 最後の単語を取得
        predictions = predictions[:, -1:, :]
        predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)

        # predicted_idがEND_TOKENと一致したら予測を終了
        if tf.equal(predicted_id, END_TOKEN[0]):
            break

        # predicted_idをこれまでのoutputの最後に追加し、次のoutputに使う
        output = tf.concat([output, predicted_id], axis=-1)

    return tf.squeeze(output, axis=0)

上記の関数を使ってcallback関数を定義、以下sampleに対する応答を各epoch終了時に返すようにする。

class Pred(tf.keras.callbacks.Callback):
    def __init__(self, start_token, end_token, max_length, validation_data=(), interval=10):
        super().__init__()
        self.START_TOKEN = start_token
        self.END_TOKEN = end_token
        self.MAX_LENGTH = max_length
        
    def on_epoch_end(self, epoch, logs={}):
        sample = [
            'こんにちは',
            '元気ですか',
            '今日の天気はどんな感じ'
        ]
        for sentence in sample:
            self._pred(sentence)

    def _pred(self, sentence):
        prediction = evaluate(sentence, tokenizer, self.model, self.START_TOKEN, self.END_TOKEN, self.MAX_LENGTH)

        predicted_sentence = tokenizer.decode(
          [i for i in prediction if i < tokenizer.vocab_size])

        print('Input: {}'.format(sentence))
        print('Output: {}'.format(predicted_sentence))

学習

epoch数とcallbackを定義して、学習を実行します。 callbackは以下の2つを指定します。

  1. チェックポイントの保存
  2. 上述したPred classでepoch終了時にsample予測

まず、チェックポイントを保存するディレクトリを作成します。 学習時間ごとにディレクトリを作成して、その中にチェックポイントを保存するようにします。

from datetime import datetime

now = datetime.now().strftime("%Y%m%d-%H%M%S")

path = "/checkpoints/" + now + "/"
import os
os.mkdir(path)

保存場所の指定と作成が終わったのでepoch数の指定と、callbackの指定をしていきます。

EPOCHS = 20


callbacks = [
    tf.keras.callbacks.ModelCheckpoint(path, monitor='loss', verbose=0, save_best_only=True, save_weights_only=True, mode='auto'),
    Pred(START_TOKEN, END_TOKEN, MAX_LENGTH)
]

以上で準備完了です。 ようやく学習開始できます。

history = model.fit(dataset, epochs=EPOCHS, callbacks=callbacks)

以下のようなログが出力されます。

Epoch 1/20
916/917 [============================>.] - ETA: 0s - loss: 1.6449 - accuracy: 0.0357
Input: こんにちは
Output: うん
Input: 元気ですか
Output: うん
Input: 今日の天気はどんな感じ
Output: うん

...

Epoch 6/20
916/917 [============================>.] - ETA: 0s - loss: 0.9231 - accuracy: 0.0934
Input: こんにちは
Output: はじめまして
Input: 元気ですか
Output: そう 
Input: 今日の天気はどんな感じ
Output: じゃあ 

...

Epoch 20/20
916/917 [============================>.] - ETA: 0s - loss: 0.6053 - accuracy: 0.1219
Input: こんにちは
Output: はじめまして
Input: 元気ですか
Output: じゃあ 
Input: 今日の天気はどんな感じ
Output: じゃあ

返答がかわっていってるのがわかって面白いです。 ですが、返答内容はバリエーションも少なく、うまくいっているとはとても言い難いです。 データセットも小さく、ハイパーパラメーターのチューニングも一切行っていないので、まだまだ改良の余地はあると思います。 いずれ再挑戦したいと思います!!!