Keras + RNN の文章生成チュートリアルをリファクタしてみた

• 概要
• 環境
• サンプルコード
• ちょっと解説
• リファクタリング不足点
• 最後に
• 参考サイト

概要

過去にやった文章生成のチュートリアルをリファクタして更に理解を深めてみました

環境

• macOS 14.0
• Python 3.11.6
  • keras 2.14.0
  • tensorflow 2.14.0
  • matplotlib 3.8.1

サンプルコード

import os
from dataclasses import dataclass
from typing import Optional

import numpy as np
import tensorflow as tf
from keras.src.callbacks import History
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.layers import GRU, Dense, Embedding
from tensorflow.keras.losses import sparse_categorical_crossentropy
from tensorflow.keras.models import Sequential
from tensorflow.keras.utils import get_file


# テストデータを管理するクラス
class TestData:
    # 文章を分割する長さ
    SEQ_LENGTH = 100
    # シャッフル用のバッファサイズ
    BUFFER_SIZE = 10000
    # シャッフル後の分割単位
    BATCH_SIZE = 64

    def __init__(self) -> None:
        # テキストデータのダウンロード
        self.__download()
        # 各種文字と文字インデックスの組み合わせの辞書と配列を作成
        self.char2idx = self.__gen_char2idx()
        self.idx2char = self.__gen_idx2char()
        # テキスト情報を数字情報のインデックス番号の配列に変換
        self.text_as_int = np.array([self.char2idx[c] for c in self.text])
        # TensorSliceDataset の作成、配列で要素は tf.Tensor に変換されたテキスト情報を数字情報のインデックス番号がある
        self.char_dataset = tf.data.Dataset.from_tensor_slices(self.text_as_int)
        # 100 文字ごとに分割して TensorSliceDataset を再度作成、BatchDataset になる
        # ここで説明変数と目的変数が生成される
        self.sequences = self.char_dataset.batch(
            self.SEQ_LENGTH + 1, drop_remainder=True
        )
        # 各 100 文字ごとの文章情報を1文字づらした [input: tf.Tensor, target: tf.Tensor] 配列を作ります、全文章分 map で繰り返して作ります
        dataset = self.sequences.map(self.shift)
        # シャルフルし再度 64 個ごとにまとめます、BatchDataset になる
        self.dataset = dataset.shuffle(self.BUFFER_SIZE).batch(
            self.BATCH_SIZE, drop_remainder=True
        )

    @property
    def vocab(self) -> list:
        # 文字を分割して辞書順に並び替えたリストを生成
        return sorted(set(self.text))

    def __download(self):
        path_to_file = get_file(
            "shakespeare.txt",
            "https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt",
        )
        self.text = open(path_to_file, "rb").read().decode(encoding="utf-8")

    def __gen_char2idx(self) -> dict:
        # 文字に対応した数字のインデックス番号の組み合わせを生成
        return {u: i for i, u in enumerate(self.vocab)}

    def __gen_idx2char(self) -> np.ndarray:
        # 文字を分割して辞書順に並び替えたリストを ndarray 形式に変換
        return np.array(self.vocab)

    def shift(self, chunk: tf.raw_ops.BatchDataset):
        # RNN 学習のために1文字づらしたシーケンスを作成します
        # :-1 は最後から1文字抜いた先頭からすべての文字 (説明変数)
        independent = chunk[:-1]
        # 1: は先頭から1文字抜いた最後までのすべての文字 (目的変数)
        dependent = chunk[1:]
        return independent, dependent


# RNN モデルを管理するクラス
class RNNModel:
    EMBEDDING_DIM = 256  # 埋め込み層の次元数
    RUN_UNITS = 1024  # RNN ユニット数
    model: Sequential

    def __init__(self, vocab_size: int, batch_size) -> None:
        self.model = Sequential()
        self.vocab_size = vocab_size
        self.batch_size = batch_size

    def build(self, batch_size: Optional[int] = None, new: bool = False):
        if batch_size is None:
            bs = self.batch_size
        else:
            bs = batch_size
        if new:
            self.model = Sequential()
        # 埋め込み層、固定次元の密ベクトルに変換するレイヤー
        self.model.add(
            Embedding(
                self.vocab_size,
                self.EMBEDDING_DIM,
                batch_input_shape=[bs, None],
            )
        )
        # RNN 層、今回は GRU レイヤーを使う
        self.model.add(
            GRU(
                self.RUN_UNITS,
                return_sequences=True,
                stateful=True,
                recurrent_initializer="glorot_uniform",
            )
        )
        # 出力層
        self.model.add(Dense(self.vocab_size))

    def compile(self) -> None:
        self.model.compile(optimizer="adam", loss=Loss().keras_scc)

    def show(self):
        self.model.summary()

    def train(
        self,
        dataset: tf.raw_ops.BatchDataset,
        epochs=10,
        callbacks=[],
    ) -> History:
        history = self.model.fit(dataset, epochs=epochs, callbacks=callbacks)
        return history

    def save(self, file_name="my_model"):
        self.model.save(file_name)

    def rebuild(self, checkpoint_dir: str = "./training_checkpoints"):
        # テスト時の入力のバッチサイズは1になるので batch_size=1 で再度ビルドする
        self.build(batch_size=1, new=True)
        self.model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
        self.model.build(tf.TensorShape([1, None]))


# 損失関数を管理するクラス
class Loss:
    def keras_scc(self, y_true, y_pred):
        # sparse_categorical_crossentropy をカスタム
        return sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)


# モデルを再度構築するので1回目のモデルを保存するためのコールバック
class Callback:
    @classmethod
    def save_model(cls) -> ModelCheckpoint:
        checkpoint_dir = "./training_checkpoints"
        checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")
        return ModelCheckpoint(filepath=checkpoint_prefix, save_weights_only=True)


# 実際にモデルを評価するためのクラス
@dataclass
class Result:
    text: str

    def show(self):
        print(self.text)


class Test:
    # 1.0 より大きいほど意外なテキストを生成する、1.0 は予測したそのままの値になる
    TEMPERRATURE = 1.0
    # 生成する文字数
    NUM_GENERATE_CHAR = 1000

    def __init__(
        self, model: RNNModel, test_data: TestData, start_string: str = "ROMEO: "
    ) -> None:
        self.model = model
        self.test_data = test_data
        self.start_string = start_string
        # start_string を数字のインデックス配列に変換
        input_eval = [self.test_data.char2idx[char] for char in start_string]
        self.input_eval = tf.expand_dims(input_eval, 0)

    def run(self) -> Result:
        self.model.model.reset_states()
        generated_text = []
        for _ in range(self.NUM_GENERATE_CHAR):
            # 文字の予測
            predictions = self.model.model(self.input_eval)
            # サイズが1の次元を削除し次元数を減らす
            predictions = tf.squeeze(predictions, 0)
            # カテゴリー分布をつかってモデルから返された文字のインデックス番号を予測
            predictions = predictions / self.TEMPERRATURE
            # 最後の要素が入力に対して予測された文字のインデックス番号
            predicted_id = tf.random.categorical(predictions, num_samples=1)[
                -1, 0
            ].numpy()
            # 過去の隠れ状態とともに予測された文字をモデルへのつぎの入力として渡す
            self.input_eval = tf.expand_dims([predicted_id], 0)
            # インデックス番号から文字への変換と結果の格納
            generated_text.append(self.test_data.idx2char[predicted_id])
        return Result(text=self.start_string + "".join(generated_text))


if __name__ == "__main__":
    # テストデータ作成
    test_data = TestData()
    # モデル生成
    model = RNNModel(len(test_data.vocab), TestData.BATCH_SIZE)
    model.build()
    model.compile()
    # # モデルの訓練
    model.train(test_data.dataset, callbacks=[Callback.save_model()])
    # モデルの評価
    model.rebuild()
    model.show()
    test = Test(model, test_data)
    result = test.run()
    # 結果を表示
    result.show()

ちょっと解説

テストデータは単純なテキストを使います
テキスト情報を1文字ずつ分解し数字とのマッピング情報を作成しテキストをすべて数値化することで学習データを作成します
文字情報の数値化は自然言語ではよくある手法になります

データをシャッフルしたりバッチサイズで分割するのは Keras で学習させる際のフォーマットに変換していると理解していますがこの辺りの手法を自然に使いこなせるレベルにならないとダメかなと思います

モデルの生成は Embedding -> GRU -> Dense とシンプルです
Embedding は自然言語の学習ではほぼ必須です
GRU (Gated Recurrent Unit) は RNN の一種です
モデル学習時にチェックポイントを保存しています
評価する際に入力をシンプルにしたいので学習時のバッチサイズ64を再度チェックポイントからモデルをビルドしてバッチサイズを1にするためです

予測した値はそのまま使わずカテゴリー分布という手法を使って再度計算させています (正直理由は不明です

リファクタリング不足点

• TestData をもう少しリファクタリングしたほうがいい
  • データの生成を init ではなくそれぞれの関数にしたほうがいい
  • 説明変数と目的変数を別のクラスで管理した方がいい
• Model の model の管理をリファクタリングしたほうがいい
  • build で新規 or 既存の書き換えを制御するよりかは ModelFactory を作って model を個別に生成できるような仕組みにするといいかも
• Test を柔軟にしたほうがいい
  • いろいろなデータでテストできるようにしたほうがいい
  • 生成文字列サイズなども可変にできるといい
  • 何をやっているかわからないところがあるのでコメントで補完したい

最後に

Keras を使った学習方法や流れはだいたい把握できましたがフルスクラッチでゼロから numpy や keras を使ってモデルを作るにはまだまだ学習が足りないと感じました

numpy で生成したベクトルの扱い方や keras のレイヤーの生成方法やレイヤーに対する入出力の制御、最適化あたりを駆使できるようにならないとオリジナルなモデルを作るのは難しいなと感じました

参考サイト

• https://keras.io/ja/
• https://www.tensorflow.org/api_docs/python/tf/raw_ops/BatchDataset
• https://www.tensorflow.org/tutorials/text/text_generation?hl=ja