0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

【その4】【自然言語処理】単語はベクトルではなく“非線形関数” ― 複素テイラー単語関数言語モデル ― 意味を“観測可能な量子状態”として扱うNLPモデルを作ってみた

0
Last updated at Posted at 2026-06-22

この記事はシリーズものの4回目です。

1回目の記事はこちらです。
2回目の記事はこちらです。
3回目の記事はこちらです。

5回目の記事はこちらです。
6回目の記事はこちらです。

↓1回目の記事:

↓2回目の記事:

↓前回(3回目)の記事:

↓次回(5回目)の記事:

↓6回目の記事:

0. 前回までのあらすじ

$$\huge{「意味」とは「状態」だ。}$$

$$\huge{「単語」とは状態を変換する「非線形作用素」だ。}$$

こう考えると、「文」は「作用素の合成」として理解することが出来るし、
「意味を理解すること」とは「状態の観測」として説明することが出来る。

image.png

image.png

―こう考えて、量子力学の真似事をしながら、新しい自然言語処理モデル「複素テイラー単語関数言語モデル」の理論を組み立てたのが前々回までだ。

そして、コーディングに着手しはじめ、モデルクラス ComplexTaylorWordFunctionLM.py を書いたのが前回だ。

前回までで、4-4節まで書いた。
今回の記事は、4-5節から開始する。

4. コーディング

複素テイラー単語関数言語モデルを学習させて実験してみるためのコードを、PyTorch で作ってみよう。

ディレクトリ構成とステータス
(カレントディレクトリ)
├ main.py                               # 😴メインプログラム
├ ComplexTaylorWordFunctionLM.py        # ✅モデルクラス
├ ComplexTaylorWordFunctionLearner.py   # 🛠️モデルを学習するためのクラス
├ constants.py                          # ✅複数ファイルで共有する定数など (コーパス以外)
├ corpus.py                             # 😴トイコーパス
├ ComplexTaylorWordFunctionEvaluator.py # 😴学習結果や経過を評価するためのクラス
└ test_project.py                       # 📌テスト用コード

ステータスの凡例
😴: 未着手
🛠️: 今から作成・編集する
📌: 新たに編集が必要になる
👀: 作成済み(未テスト)
✅: 作成済み (テスト成功済み or テストを実施しない)
⚠️: 問題あり

まずは 4-5節~4-7説で、学習器クラス ComplexTaylorWordFunctionLearner.py をコーディングする。
次に、4-8 節で学習器クラスをテストする。
その後、次回以降の記事でその他のファイルをコーディングしていく。

4-5. 学習器クラスの設計

ComplexTaylorWordFunctionLM モデルを学習する ComplexTaylorWordFunctionLearner クラスを作ろう。

親クラスは指定しない (=object が親になる)。

まずは:

  • 属性(メソッド, データ属性)概要
  • 訓練データ
  • ハイパーパラメータ
  • オプティマイザ

を決めよう。

4-5-1. メソッド一覧

ComplexTaylorWordFunctionLearner クラスには次のメソッドを持たせる。
(すべてインスタンスメソッドである。)

メソッド 役割
__init__(コンストラクタ) モデル、訓練データ、オプティマイザを初期化する
set_optimizer オプティマイザクラスやハイパーパラメータを受け取って、オプティマイザを初期化する。
初期化されたオプティマイザを返す。
add_words 単語のリストを受け取り、それらの単語をモデルの語彙に追加する。
追加に成功した単語のリストを返す。
remove_words 単語のリストを受け取り、それらの単語をモデルの語彙から削除する。
削除に成功した単語のリストを返す。
make_examples 文のリストを受け取り、訓練データ examples を返す。
データ属性 learner.examples は更新しないことに注意
make_batchs バッチサイズを受け取り、訓練データ learner.examples をバッチのリストに分けたものを返す。
make_batches (make_batchs のエイリアス)
make_batch_tensor バッチを受け取り、単語IDのテンソルのタプルに変換して返す。
オプション引数によっては、未知語を検知した場合にそれらをモデルの語彙に追加する。
update_cache モデルのキャッシュを更新する。何も受け取らず、何も返さない。
operator_l2_regularization 何も受け取らず、モデルの単語作用素のパラメータの絶対値の二乗の平均を返す
normalize_batch バッチを受け取り、所定のフォーマットに変換したものを返す
learn バッチを受け取り、$1$ステップ学習する。
ステップ結果の情報を返す。
zero_grad 勾配情報を消す。何も受け取らず、何も返さない。
accumulate_grad バッチを受け取り、勾配を蓄積する。
そこで得られる損失関数の値などの情報を返す。
on_before_word_added フック。モデルの語彙に単語が追加される直前に発火する。
入力は追加される単語のリストで、出力は無視される。
on_word_added フック。モデルの語彙に単語が追加された直後に発火する。
入力は追加された(単語とそれに与えられた単語IDのタプル)のリストで、出力は無視される。
on_strange_ignored フック。learn()メソッドや make_batch_padded()メソッドで未知語を検出し、それが無視される (strange="ignore") ときに発火する。
入力は無視される単語のリストで、出力は無視される。
on_before_word_removed フック。モデルの語彙から単語が削除される直前に発火する。
入力は削除される単語のリストで、出力は無視される。
on_word_removed フック。モデルの語彙から単語が削除された直後に発火する。
入力は削除された単語のリストで、出力は無視される。

実際に学習を行う 最重要メソッドは learn である。
learn メソッドに引数として渡される訓練データのバッチは make_batchs メソッドで作られるので、 make_batchs メソッドも重要である。
そして 訓練データのバッチは、 make_examples メソッドで作られる訓練データ examples から作られるので、 make_examples メソッドも重要である。

上記の重要3メソッドをまとめると:

  1. make_examples: 文のリストから訓練データ examples を作る
  2. make_batchs: examples から訓練データのバッチ batch を作る
  3. learn: batch を使って学習する

といった関係である。

4-5-2. 訓練データ

訓練データは、文 (list[str]) の両端に null トークンを付けたものから、「一定の長さ」以下だけ切り出した 問題 (list[str]) と
その直後に来る単語 である正解 (str) を、
タプルにまとめたリスト:

訓練データ
list[              # 訓練データ
    tuple[         # 訓練データム (データム: データの単数形)
        list[str], # 問題
        str        # 正解
    ]
]

とすることにする。

そして、この「一定の長さ」のことを context_len と呼んだりする。

例えば context_len == 4 で、
["I", "love", "you", "very", "much"] という文から訓練データを作ると:

訓練データ
examples == [
    (["<Null>"],                      "I"), 
    (["<Null>", "I"],                 "love"), 
    (["<Null>", "I", "love"],         "you"), 
    (["<Null>", "I", "love", "you"],  "very"), 
    (["I", "love", "you", "very"],    "much"), 
    (["love", "you", "very", "much"], "<Null>")
]

となる。
両端に null トークンを付けたことで、文の「最初」と「最後」の概念を学習させることが出来る。

訓練データム (["<Null>"], "I") からは「何もないの次」、つまり最初の単語が"I" であることが学習できるし、
訓練データム (["love", "you", "very", "much"], "<Null>") からは
"... love you very much" の次が「何もない」、つまり "... love you very much" が最後であることを学習できる。

ここでは一つの文から 訓練データを作る場合を例示したが、実際には複数の文から訓練データを作ることが出来る。
その場合には、各文から作られる訓練データをマージして一つの大きな訓練データにすることになる。

そして、訓練データ examples (list[tuple[list[str], str]]) を、訓練データム example が 最大 (バッチサイズ) 個集められたリスト (list[tuple[list[str], str]])になるようまとめなおしたもの(list[list[tuple[list[str], str]]]) を batchs といい、その要素をバッチ batch という。

4-5-3. ハイパーパラメータ

ハイパーパラメータには、先ほど出てきた訓練データの「問題」の長さである context_len も含め、次のようなものがある。

他にも増やそうと思えばいくらでもハイパーパラメータを増やすことは出来るが、今回の実験ではとりあえずこれだけ考えよう。

名前 説明
batch_size int バッチサイズ
content_len int 問題の長さ
base_lr float 単語作用素の学習率
next_lr float $\mathbb{next}$ の学習率
weight_decay float 重み減衰
lambda_norm float L2正則化の係数
clip_grad_norm float 勾配の大きさ制限

4-5-4. オプティマイザ

学習器クラスに公開インスタンスデータ属性 opt を持たせて、実験者が自由にいじれる仕様にしたい。

(そうしないと、「単語作用素と$\mathbb{next}$で学習率を分けたい」みたいなワガママに答えるのが難しくなる。)

とはいえ、ある程度どのようなオプティマイザを使うか目星はつけておきたいので、
とりあえず torch.optim.AdamW, torch.optim.Adam, torch.optim.SGD のどれかを使う、と考えておこう。

4-5-5. データ属性一覧

ComplexTaylorWordFunctionLearner クラスは次のデータ属性を持つ。
(すべて インスタンスデータ属性である。)

変数 種類 内容
model モデル 学習させたい ComplexTaylorWordFunctionLM モデル
strange 設定 学習用バッチに未知語が含まれていた時、モデルの語彙に追加するか、無視するか、エラーを起こすか
auto_update_cache_before_learning 設定 重み更新直前に model のキャッシュを更新するかどうか
auto_update_cache_after_learning 設定 重み更新直後に model のキャッシュを更新するかどうか
norm_reg_mode 設定 L2正則化の対象を通常の単語作用素のみにするか、$\mathbb{next}$のみにするか、両方にするか
lambda_norm 設定 L2正則化用の係数
clip_grad_norm 設定 勾配ベクトルの大きさの制限
optimizer_rebuilt フラグ モデルの語彙が変化したときに、それに合わせてオプティマイザが同期されたか否か。
特筆事項あり (4-5-6節)
examples 訓練データ make_batchs メソッドでバッチにまとめられる訓練データ
opt オプティマイザ learn メソッドで使われるオプティマイザ

ちなみに:

  • 問題の最大長 (context_len) については、訓練データ (examples) を作る際に必要なだけなので、データ属性として記録する必要はない。
  • バッチサイズについては、 make_batchs メソッドでバッチを作る際に毎回指定する仕様でいいだろうから、データ属性として記録しない。
  • 学習率 (lr) と 重み減衰 (weight_decay) については、オプティマイザに記録されるので、学習器クラスがデータ属性として記録する必要はない。

4-5-6. optimizer_rebuilt データ属性

optimizer_rebuilt 属性は、オプティマイザとモデルの間に矛盾がないかを示すフラグとして使われる。

今回の学習器は add_wordsremove_words メソッドによりモデルの語彙を増やしたり減らしたりできる仕様になっているが、これをやった場合、学習パラメータが増減することになる。
従って、オプティマイザを作り直さないと以降正しく学習を出来なくなる。

そこで、add_wordsremove_words メソッドが呼ばれると
optimizer_rebuilt 属性 が False になり、
set_optimizer メソッドが呼ばれると
optimizer_rebuilt 属性 が True になる仕様にする。

そして learn メソッドや accumulate_grad メソッド が、 learner.optimizer_rebuilt == False
の時に呼ばれるとエラーを投げる仕様にすれば、誤った学習を防ぐことが出来る。

4-6. 学習器クラスの各メソッドのコーディング

フックメソッド以外の各メソッドをコーディングしていく。

4-6-1. コンストラクタ

コンストラクタがやるべきことは、結局は4-5-5節で列挙したデータ属性の初期化に過ぎない。


まずは引数。

コンストラクタの引数
def __init__(
    self, sentences, d=16, k=4, next_identity_init=True,
    strange="error", optimizer="AdamW",
    lr=7e-4, weight_decay=0.0, lambda_norm=1.0, clip_grad_norm=1.0,
    context_len=8, auto_update_cache="after", norm_reg_mode="both",
    device=None
):

sentences には文のリスト (list[list[str]]) を指定する。
コンストラクタは sentences から 語彙 vocab を組み立て、それを用いて内部でモデルself.model を初期化する。
d, k, next_identity_init もまた、 self.model の初期化に使われる。

strange には、learn メソッドに渡された学習用バッチに未知語が含まれていた場合のポリシーを指定する。

  • strange == "add": モデルの語彙に未知語を追加する
  • strange == "ignore": 未知語を "<Null>" で置き換える
  • strange == "error": 未知語を検出したときに ValueError を投げる

optimizer には 初期化の時点で採用したいオプティマイザの名前 ("AdamW", "Adam", "SGD" から選ぶ)か、 torch.optim.Optimizer継承クラス、または None を指定する。
lrweight_decay には、 オプティマイザに指定したい学習率と重み減衰を指定する。

lambda_norm には、L2正則化用の係数を指定する。

clip_grad_norm には、勾配ベクトルの大きさの制限を指定する。
(制限しない場合は False を指定する)

context_len には、問題の最大長を指定する。

auto_update_cache には、学習の前後に モデルのキャッシュを更新するかを指定する。

  • auto_update_cache == "before": 学習前に更新する
  • auto_update_cache == "after": 学習後に更新する
  • auto_update_cache in ["both", True]: 学習前後両方に更新を行う
  • auto_update_cache in ["none", False]: 更新しない

norm_reg_mode には、L2正則化の際に見るモデルパラメータを指定する。

  • norm_reg_mode == "A": 通常の単語作用素パラメータをL2正則化する
  • norm_reg_mode == "Next": $\mathbb{next}$パラメータをL2正則化する
  • norm_reg_mode == "both": 通常の単語作用素と$\mathbb{next}$ の両方のパラメータをL2正則化する

device には、モデルを割り当てたい場所 (CPU / GPU など)を指定する。
こだわりがないなら None のままでよい。


次に「設定」と「フラグ」のデータ属性を初期化しよう。

次のようにすればよい。

コンストラクタで「設定」用と「フラグ」のデータ属性を初期化
def __init__(
    self, sentences, d=16, k=4, next_identity_init=True,
    strange="error", optimizer="AdamW",
    lr=7e-4, weight_decay=0.0, lambda_norm=1.0, clip_grad_norm=1.0,
    context_len=8, auto_update_cache="after", norm_reg_mode="both",
    device=None
):
    strange = str(strange).lower()
    if strange not in {"add", "ignore", "error"}:
        raise ValueError('strange must be "add", "ignore", or "error"')
    self.strange = strange
    
    err = ValueError("auto_update_cache argument is illegal.")
    if   isinstance(auto_update_cache, bool):
        before = bool(auto_update_cache)
        after  = bool(auto_update_cache)
    elif isinstance(auto_update_cache, str):
        auto_update_cache = str(auto_update_cache).lower()
        if auto_update_cache not in {"before", "after", "both", "none"}:
            raise err
        before = (auto_update_cache in {"both", "before"})
        after  = (auto_update_cache in {"both", "after"})
    else:
        raise err

    self.auto_update_cache_before_learning = before
    self.auto_update_cache_after_learning = after

    norm_reg_mode = str(norm_reg_mode).lower()        
    if norm_reg_mode not in {"a", "next", "both"}:
        raise ValueError("norm_reg_mode argument is illegal.")
    self.norm_reg_mode = norm_reg_mode
    
    self.lambda_norm = lambda_norm
    
    if isinstance(clip_grad_norm, bool) and clip_grad_norm:
        clip_grad_norm = 1.0
    if clip_grad_norm is None:
        clip_grad_norm = False
    self.clip_grad_norm = clip_grad_norm
            
    self.optimizer_rebuilt = False
    ...

この時点でやっていることは:

  • 引数 strange から データ属性 strange を初期化する

  • 引数 auto_update_cache から データ属性:

    • auto_update_cache_before_learning
    • auto_update_cache_after_learning

    を初期化する

  • 引数 norm_reg_mode から データ属性 norm_reg_mode を初期化する

  • 引数 lambda_norm から データ属性 lambda_norm を初期化する

  • 引数 clip_grad_norm から データ属性 clip_grad_norm を初期化する

  • データ属性 optimizer_rebuiltFalse で初期化する (オプティマイザをまだ作っていないので)

ことである。


次に、モデルを初期化しよう。

次のようにすればよい。

コンストラクタで「モデル」のデータ属性を初期化
def __init__(
    self, sentences, d=16, k=4, next_identity_init=True,
    strange="error", optimizer="AdamW",
    lr=7e-4, weight_decay=0.0, lambda_norm=1.0, clip_grad_norm=1.0,
    context_len=8, auto_update_cache="after", norm_reg_mode="both",
    device=None
):
    ...
    # ここから追加
    words = []
    seen = set()
    for sentence in sentences:
        for word in sentence:
            if word != NULL_TOKEN and word not in seen:
                seen.add(word)
                words.append(word)

    vocab = sorted(words)
    vocab.insert(NULL_ID, NULL_TOKEN)

    self.model = ComplexTaylorWordFunctionLM(
        vocab, d=d, k=k, next_identity_init=next_identity_init,
    )

    if device is not None:
        self.model.to(torch.device(device))
    self.model.reset_null()
    ...

まず、 wordssentences に出現する単語を全て抽出している。
そしてそれを vocab と呼び変え、
vocab に対して、決められたインデックス(NULL_ID)に "<Null>" (NULL_TOKEN) を追加している。

この vocab を使って モデル self.model を初期化している。

そして引数 device が指定されている場合は self.model.to で モデルを当該デバイスへ移動している。

最後の self.model.reset_null() は、モデル初期化時にも呼ばれているので、本当は要らない。


次に訓練データとオプティマイザを初期化しよう。

次のようにすればよい。

コンストラクタで「訓練データ」と「オプティマイザ」のデータ属性を初期化
def __init__(
    self, sentences, d=16, k=4, next_identity_init=True,
    strange="error", optimizer="AdamW",
    lr=7e-4, weight_decay=0.0, lambda_norm=1.0, clip_grad_norm=1.0,
    context_len=8, auto_update_cache="after", norm_reg_mode="both",
    device=None
):
    ...
    # ここから追加
    self.examples = []
    self.examples += self.make_examples(sentences, context_len)

    self.opt = None
    self.set_optimizer(
        optimizer=optimizer,
        lr=lr,
        weight_decay=weight_decay,
    )

訓練データは make_examples メソッドで作成されたものを、空リストのexamples 属性へ追加登録しているに過ぎない。

またオプティマイザは set_optimizer メソッドで登録しており、このメソッドで opt 属性が更新される。


コンストラクタ全体のコードは次のとおり。

コンストラクタ全体のコードを表示
ComplexTaylorWordFunctionLearnerクラスのコンストラクタ
def __init__(
    self, sentences, d=16, k=4, next_identity_init=True,
    strange="error", optimizer="AdamW",
    lr=7e-4, weight_decay=0.0, lambda_norm=1.0, clip_grad_norm=1.0,
    context_len=8, auto_update_cache="after", norm_reg_mode="both",
    device=None
):
    # 1. 設定, フラグ
    strange = str(strange).lower()
    if strange not in {"add", "ignore", "error"}:
        raise ValueError('strange must be "add", "ignore", or "error"')
    self.strange = strange
    
    err = ValueError("auto_update_cache argument is illegal.")
    if   isinstance(auto_update_cache, bool):
        before = bool(auto_update_cache)
        after  = bool(auto_update_cache)
    elif isinstance(auto_update_cache, str):
        auto_update_cache = str(auto_update_cache).lower()
        if auto_update_cache not in {"before", "after", "both", "none"}:
            raise err
        before = (auto_update_cache in {"both", "before"})
        after  = (auto_update_cache in {"both", "after"})
    else:
        raise err

    self.auto_update_cache_before_learning = before
    self.auto_update_cache_after_learning = after

    norm_reg_mode = str(norm_reg_mode).lower()        
    if norm_reg_mode not in {"a", "next", "both"}:
        raise ValueError("norm_reg_mode argument is illegal.")
    self.norm_reg_mode = norm_reg_mode
    
    self.lambda_norm = lambda_norm
    
    if isinstance(clip_grad_norm, bool) and clip_grad_norm:
        clip_grad_norm = 1.0
    if clip_grad_norm is None:
        clip_grad_norm = False
    self.clip_grad_norm = clip_grad_norm
            
    self.optimizer_rebuilt = False

    # 2. モデル
    words = []
    seen = set()
    for sentence in sentences:
        for word in sentence:
            if word != NULL_TOKEN and word not in seen:
                seen.add(word)
                words.append(word)

    vocab = sorted(words)
    vocab.insert(NULL_ID, NULL_TOKEN)

    self.model = ComplexTaylorWordFunctionLM(
        vocab, d=d, k=k, next_identity_init=next_identity_init,
    )

    if device is not None:
        self.model.to(torch.device(device))
    self.model.reset_null()

    # 3. 訓練データ
    self.examples = []
    self.examples += self.make_examples(sentences, context_len)

    # 4. オプティマイザ
    self.opt = None
    self.set_optimizer(
        optimizer=optimizer,
        lr=lr,
        weight_decay=weight_decay,
    )

(コンストラクタ全体のコードを表示 ここまで)

4-6-2. set_optimizer メソッド

set_optimizer メソッドでは、オプティマイザの種類やオプティマイザ用のハイパーパラメータを受け取り、
オプティマイザを初期化してデータ属性 opt を更新する。
optimizer_rebuilt 属性を Trueにする。
そして opt をそのまま返す。


コードは次のとおり。

ComplexTaylorWordFunctionLearnerクラスのset_optimizerメソッド
def set_optimizer(
    self, optimizer="AdamW", lr=7e-4, weight_decay=0.0, **kwargs
):
    if optimizer is None:
        self.opt = None
        return None

    if isinstance(optimizer, str):
        name = optimizer.lower()
        if name == "adamw":
            opt_cls = torch.optim.AdamW
        elif name == "adam":
            opt_cls = torch.optim.Adam
        elif name == "sgd":
            opt_cls = torch.optim.SGD
        else:
            raise ValueError(f"unknown optimizer: {optimizer}")
    else:
        opt_cls = optimizer

    self.opt = opt_cls(
        [p for p in self.model.parameters() if p.requires_grad],
        lr=lr,
        weight_decay=weight_decay,
        **kwargs,
    )
    self.optimizer_rebuilt = True
    return self.opt

引数 optimizerNone の場合は、学習器からオプティマイザを取り除いている。
引数 optimizerstr 型の場合は、"AdamW", "Adam", "SGD" から選ぶことになる。
引数 optimizerNoneでもstr型でもない場合、
torch.optim.Optimizer 継承型を表す type オブジェクトであるとみなされ、それに対応するオプティマイザが採用される。

そして最後にハイパーパラメータをすべてオプティマイザに渡し、
これを opt 属性に登録し、
optimizer_rebuilt 属性を True にして
opt 属性を返している。

4-6-3. add_words メソッド

add_words メソッド は、単語のリストを受け取り、それらの単語をモデルの語彙に追加する。
追加に成功した単語のリストを返す。
optimizer_rebuilt 属性を変更する。
フック on_before_word_added を発火する。
フック on_word_added を発火する


まず、optimizer_rebuilt属性をFalseにしよう。
単語をモデルの語彙に追加するということは、オプティマイザが古くなるということだから。

add_wordsメソッドでoptimizer_rebuilt属性をFalseにする
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    self.optimizer_rebuilt = False
    ...

次に、本当に未知語である単語だけ抽出しよう。
もし words の中に、既にモデルの語彙に存在するものがあるなら、
それは除いてよいからだ。

add_wordsメソッドで未知語を抽出する
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    ...
    # ここから追加
    if isinstance(words, str):
        words = [words]
    
    stoi = dict(self.model.stoi)
    itos = dict(self.model.itos)
    vocab = list(itos.values())
    new_words = []
    for w in words:
        if w not in stoi:
            stoi[w] = len(vocab)
            itos[len(vocab)] = w
            vocab.append(w)
            new_words.append(w)
    ...

new_words が未知語のリストである。


次に フック on_before_word_added を発火しよう

add_wordsメソッドでフックを発火する
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    ...
    # ここから追加
    self.on_before_word_added(new_words)
    ...

次に モデルに未知語を追加しよう

add_wordsメソッドでモデルに未知語を追加する
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    ...
    # ここから追加
    old_V = self.model.V
    add_V = len(new_words)
    new_V = old_V + add_V

    self.model.V = new_V
    self.model.stoi = dict(stoi)
    self.model.itos = dict(itos)

    for n in range(self.model.k + 1):
        old_A = self.model.A[n].detach()
        new_A_rows = MAKE_A(add_V, self.model.d, n).to(
            dtype=old_A.dtype, device=old_A.device
        )
        expanded_A = torch.cat([old_A, new_A_rows], dim=0)
        self.model.A[n] = nn.Parameter(expanded_A)

    old_meanings = self.model.word_meanings_povm.detach()
    extra_meanings = torch.zeros(
        add_V,
        self.model.d,
        dtype=old_meanings.dtype,
        device=old_meanings.device,
    )
    self.model.word_meanings_povm = torch.cat(
        [old_meanings, extra_meanings], dim=0
    )

    self.model.reset_null()
    ...

old_meanings = ... より前では、
モデル self.modelV, stoi, itos, A を新しいものに書き換えている。

old_meanings = ... から self.model.word_meanings_povm = ... では、
モデルの word_meanings_povm 属性を拡大している。
(拡大しておかないと、形状不一致のエラーが出てしまう)

また、不要だが念のためモデルを reset_null している。


次に 自動的に行う処理をしよう

add_wordsメソッドで自動的に行う処理をする
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    ...
    # ここから追加
    if rebuild_optimizer:
        self.set_optimizer(
            optimizer=self.opt.__class__,
            lr=self.opt.defaults['lr'],
            weight_decay=self.opt.defaults['weight_decay'],
        )
    if update_cache:
        self.update_cache()
    ...

引数 rebuild_optimizerTrue の場合、
set_optimizer メソッドが呼ばれるので、
属性 optimizer_rebuilt は結局 True になる。

また引数 update_cacheTrue の場合、モデルのキャッシュが自動で更新される。


次に フック on_word_added を発火しよう

add_wordsメソッドでフックを発火する
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    ...
    # ここから追加
    new_ids = [stoi[w] for w in new_words]
    self.on_word_added(list(zip(new_words, new_ids)))
    ...

最後に 未知語のリストを返そう。

add_wordsメソッドで未知語のリストを返す
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    ...
    # ここから追加
    return new_words

add_words メソッド全体のコードは次のとおり。

add_words メソッド全体のコードを表示
ComplexTaylorWordFunctionLearnerクラスのadd_wordsメソッド
def add_words(self, words, rebuild_optimizer=True, update_cache=True):
    # 1. フラグを立てる
    self.optimizer_rebuilt = False

    # 2. 未知語を抽出する
    if isinstance(words, str):
        words = [words]
    
    stoi = dict(self.model.stoi)
    itos = dict(self.model.itos)
    vocab = list(itos.values())
    new_words = []
    for w in words:
        if w not in stoi:
            stoi[w] = len(vocab)
            itos[len(vocab)] = w
            vocab.append(w)
            new_words.append(w)

    # 3. フック on_before_word_added を発火する
    self.on_before_word_added(new_words)

    # 4. モデルに未知語を追加する
    old_V = self.model.V
    add_V = len(new_words)
    new_V = old_V + add_V

    self.model.V = new_V
    self.model.stoi = dict(stoi)
    self.model.itos = dict(itos)

    for n in range(self.model.k + 1):
        old_A = self.model.A[n].detach()
        new_A_rows = MAKE_A(add_V, self.model.d, n).to(
            dtype=old_A.dtype, device=old_A.device
        )
        expanded_A = torch.cat([old_A, new_A_rows], dim=0)
        self.model.A[n] = nn.Parameter(expanded_A)

    old_meanings = self.model.word_meanings_povm.detach()
    extra_meanings = torch.zeros(
        add_V,
        self.model.d,
        dtype=old_meanings.dtype,
        device=old_meanings.device,
    )
    self.model.word_meanings_povm = torch.cat(
        [old_meanings, extra_meanings], dim=0
    )

    self.model.reset_null()

    # 5. 自動処理をする
    if rebuild_optimizer:
        self.set_optimizer(
            optimizer=self.opt.__class__,
            lr=self.opt.defaults['lr'],
            weight_decay=self.opt.defaults['weight_decay'],
        )
    if update_cache:
        self.update_cache()

    # 6. フック on_word_added を発火する
    new_ids = [stoi[w] for w in new_words]
    self.on_word_added(list(zip(new_words, new_ids)))

    # 7. 未知語のリストを返す
    return new_words

(add_words メソッド全体のコードを表示 ここまで)

4-6-4. remove_words メソッド

remove_words メソッド は、単語のリストを受け取り、それらの単語をモデルの語彙から削除する。
削除に成功した単語のリストを返す。
optimizer_rebuilt 属性を変更する。
フック on_before_word_removed を発火する。
フック on_word_removed を発火する


まず、optimizer_rebuilt属性をFalseにしよう。
単語をモデルの語彙から削除するということは、オプティマイザが古くなるということだから。

remove_wordsメソッドでoptimizer_rebuilt属性をFalseにする
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    self.optimizer_rebuilt = False
    ...

次に 引数に異常がないか確認しよう。

example_policy 引数は 訓練データ内にある「削除された単語」をどう扱うかの方針であり、
次の3通りを許す。

  • "drop": 削除された単語を含む訓練データムを取り除く
  • "null": 削除された単語は $\mathbb{null}$ で置き換える
  • "keep": 何もせずそのまま残す
    • "keep" を選んだ場合、そのままstrange="add" モードで学習を継続すると、
      せっかく削除した単語がまた復活し得るので注意

それ以外の値が指定された場合は ValueError を投げる。

削除したい単語のリストである words 引数に $\mathbb{null}$ が含まれていた場合は ValueError を投げる。

remove_wordsメソッドで引数に異常がないか確認する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    if isinstance(words, str):
        words = [words]

    example_policy = example_policy.lower()
    if example_policy not in {"drop", "null", "keep"}:
        raise ValueError('example_policy must be "drop", "null", or "keep"')

    removes = set(words)

    if NULL_TOKEN in removes:
        raise ValueError(f"{NULL_TOKEN} cannot be removed.")
    ...

次に、本当に削除すべき単語だけを抽出しよう。
もし words の中に、モデルの語彙にないものがあれば、それは無視してよいからだ。

remove_wordsメソッドで削除すべき単語を抽出する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    old_stoi = dict(self.model.stoi)
    old_itos = dict(self.model.itos)

    removed_words = [w for w in old_itos.values() if w in removes]
    ...

次に、削除を実行することでモデルの語彙が少なくなりすぎないか確認しよう。
今回のモデルでは、語彙数 が 意味ベクトルの次元 未満であってはならないからだ。

remove_wordsメソッドで語彙が少なくなりすぎないか確認する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    new_vocab = [
        old_itos[i]
        for i in range(self.model.V)
        if old_itos[i] not in removes
    ]

    if len(new_vocab) < self.model.d:
        raise RuntimeError(
            "Cannot remove words because the frame-normalized POVM "
            "assumes d <= V. "
            f"After removal, d={self.model.d}, V={len(new_vocab)}."
        )
    ...

次に、on_before_word_removed フックを発火しよう

remove_wordsメソッドでフックを発火する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    self.on_before_word_removed(removed_words)
    ...

次に、モデルから不要語を削除しよう

remove_wordsメソッドでモデルから不要語を削除する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    old_to_new = {}
    new_itos = {}
    new_stoi = {}

    for new_id, w in enumerate(new_vocab):
        old_id = old_stoi[w]
        old_to_new[old_id] = new_id
        new_itos[new_id] = w
        new_stoi[w] = new_id

    keep_old_ids = torch.tensor(
        [old_stoi[w] for w in new_vocab],
        dtype=torch.long,
        device=self.model.z0.device,
    )

    self.model.V = len(new_vocab)
    self.model.itos = new_itos
    self.model.stoi = new_stoi

    for n in range(self.model.k + 1):
        old_A = self.model.A[n].detach()
        new_A = old_A.index_select(0, keep_old_ids)
        self.model.A[n] = nn.Parameter(new_A)

    old_meanings = self.model.word_meanings_povm.detach()
    self.model.word_meanings_povm = old_meanings.index_select(
        0, keep_old_ids
    )

    self.model.reset_null()
    ...

old_meanings = ... より前では、
モデル self.modelV, stoi, itos, A を新しいものに書き換えている。

old_meanings = ... から self.model.word_meanings_povm = ... では、
モデルの word_meanings_povm 属性を縮小している。
(縮小しておかないと、形状不一致のエラーが出てしまう)

また、不要だが念のためモデルを reset_null している。


次に、訓練データ内の不要語を処理しよう。
「引数に異常がないか確認しよう」のところで説明した通り、
処理方針は引数 example_policy に基づく。

remove_wordsメソッドで訓練データ内の不要語を処理する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    removeds = set(removed_words)

    if example_policy == "drop":
        self.examples = [
            (ctx_words, target_word)
            for (ctx_words, target_word) in self.examples
            if (
                target_word not in removeds
                and all(w not in removeds for w in ctx_words)
            )
        ]

    elif example_policy == "null":
        self.examples = [
            (
                [
                    NULL_TOKEN if w in removeds else w
                    for w in ctx_words
                ],
                NULL_TOKEN if target_word in removeds else target_word,
            )
            for (ctx_words, target_word) in self.examples
        ]

    elif example_policy == "keep":
        pass
    ...

次に、自動で行う処理をしよう。

remove_wordsメソッドで自動で行う処理をする
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    if rebuild_optimizer:
        self.set_optimizer(
            optimizer=self.opt.__class__,
            lr=self.opt.defaults['lr'],
            weight_decay=self.opt.defaults['weight_decay'],
        )
    if update_cache:
        self.update_cache()
    ...

引数 rebuild_optimizerTrue の場合、
set_optimizer メソッドが呼ばれるので、
属性 optimizer_rebuilt は結局 True になる。

また引数 update_cacheTrue の場合、モデルのキャッシュが自動で更新される。


次に、フック on_word_removed を発火しよう。

remove_wordsメソッドでフックを発火する
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    self.on_word_removed(removed_words)
    ...

最後に、削除された単語のリストを返そう。

remove_wordsメソッドで削除された単語のリストを返す
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    ...
    # ここから追加
    return removed_words

remove_words メソッド全体のコードは次のとおり。

remove_wordsメソッド全体のコードを表示
ComplexTaylorWordFunctionLearnerクラスのremove_wordsメソッド
def remove_words(
    self, words, example_policy="drop",
    rebuild_optimizer=True, update_cache=True
):
    # 1. フラグを立てる
    self.optimizer_rebuilt = False

    # 2. 引数に異常がないか確認する
    if isinstance(words, str):
        words = [words]

    example_policy = example_policy.lower()
    if example_policy not in {"drop", "null", "keep"}:
        raise ValueError('example_policy must be "drop", "null", or "keep"')

    removes = set(words)

    if NULL_TOKEN in removes:
        raise ValueError(f"{NULL_TOKEN} cannot be removed.")

    # 3. 不要語を抽出する
    old_stoi = dict(self.model.stoi)
    old_itos = dict(self.model.itos)

    removed_words = [w for w in old_itos.values() if w in removes]

    # 4. 語彙が少なくなりすぎないか確認する
    new_vocab = [
        old_itos[i]
        for i in range(self.model.V)
        if old_itos[i] not in removes
    ]

    if len(new_vocab) < self.model.d:
        raise RuntimeError(
            "Cannot remove words because the frame-normalized POVM "
            "assumes d <= V. "
            f"After removal, d={self.model.d}, V={len(new_vocab)}."
        )

    # 5. フック on_before_word_removed を発火する
    self.on_before_word_removed(removed_words)

    # 6. モデルから不要語を削除する
    old_to_new = {}
    new_itos = {}
    new_stoi = {}

    for new_id, w in enumerate(new_vocab):
        old_id = old_stoi[w]
        old_to_new[old_id] = new_id
        new_itos[new_id] = w
        new_stoi[w] = new_id

    keep_old_ids = torch.tensor(
        [old_stoi[w] for w in new_vocab],
        dtype=torch.long,
        device=self.model.z0.device,
    )

    self.model.V = len(new_vocab)
    self.model.itos = new_itos
    self.model.stoi = new_stoi

    for n in range(self.model.k + 1):
        old_A = self.model.A[n].detach()
        new_A = old_A.index_select(0, keep_old_ids)
        self.model.A[n] = nn.Parameter(new_A)

    old_meanings = self.model.word_meanings_povm.detach()
    self.model.word_meanings_povm = old_meanings.index_select(
        0, keep_old_ids
    )

    self.model.reset_null()
    
    # 7. 訓練データ内の不要語を処理する
    removeds = set(removed_words)

    if example_policy == "drop":
        self.examples = [
            (ctx_words, target_word)
            for (ctx_words, target_word) in self.examples
            if (
                target_word not in removeds
                and all(w not in removeds for w in ctx_words)
            )
        ]

    elif example_policy == "null":
        self.examples = [
            (
                [
                    NULL_TOKEN if w in removeds else w
                    for w in ctx_words
                ],
                NULL_TOKEN if target_word in removeds else target_word,
            )
            for (ctx_words, target_word) in self.examples
        ]

    elif example_policy == "keep":
        pass

    # 8. 自動処理をする
    if rebuild_optimizer:
        self.set_optimizer(
            optimizer=self.opt.__class__,
            lr=self.opt.defaults['lr'],
            weight_decay=self.opt.defaults['weight_decay'],
        )
    if update_cache:
        self.update_cache()

    # 9. フック on_word_removed を発火する
    self.on_word_removed(removed_words)

    # 10. 不要語のリストを返す
    return removed_words

( remove_words メソッド全体のコードを表示 ここまで)

4-6-5. make_examples メソッド

make_examples メソッド は、次の2つを受け取り、訓練データ examples を返す。

  • 訓練データの元になる文 list[str] のリスト sentences: list[list[str]]
  • 問題の最大長 context_len

但し、訓練データを保持するデータ属性 learner.examples は更新しない。


訓練データを make_examples で作ったはよいが、それを実際に採用するか、
採用にするにしても、古い訓練データを消して上書きするのか、それとも単に追加するのか、
などの判断は学習器 learner の利用者に委ねられる。

このメソッドがデータ属性 learner.examples を更新しないのはそのためである。

例えば、訓練データを追加登録したい場合、利用者は:

learner.examples += learner.make_examples(sentences, context_len)

のようにすればよい。


コードは次のとおり。

ComplexTaylorWordFunctionLearnerクラスのmake_examplesメソッド
def make_examples(self, sentences, context_len):
    if sentences and isinstance(sentences[0], str):
        sentences = [sentences]
        
    sentences_without_null = [
        [w for w in s if w != NULL_TOKEN] for s in sentences
    ]
    
    sentences_with_null = [
        [NULL_TOKEN] + s + [NULL_TOKEN] for s in sentences_without_null
    ]

    examples = []

    for sentence in sentences_with_null:
        for t in range(1, len(sentence)):
            start = max(0, t - context_len)
            ctx_words = sentence[start:t]
            target_word = sentence[t]
            examples.append((ctx_words, target_word))

    return examples

sentences_without_null では、各文から $\mathbb{null}$ を取り除いている。
sentences_with_null では、そこからさらに各文の両端に $\mathbb{null}$ を追加している。
両端に null トークンを付けることのメリットは4-5-2節で解説した。

その後、4-5-2節で述べた通りの要領で、訓練データ examples を作って返している。

4-6-6. make_batchs / make_batches メソッド

make_batchs メソッドは、バッチサイズを受け取り、データ属性として記録された訓練データ examples をバッチのリストに分けたものを返す。

make_batches メソッドは、 make_batchs のエイリアス。


コードは次のとおり。

ComplexTaylorWordFunctionLearnerクラスのmake_batchs / make_batchesメソッド
def make_batchs(self, batch_size, shuffle=True):
    indices = list(range(len(self.examples)))
    if shuffle:
        random.shuffle(indices)

    batchs = []
    for start in range(0, len(indices), batch_size):
        batch_indices = indices[start:start + batch_size]
        batchs.append([self.examples[i] for i in batch_indices])
    return batchs

def make_batches(self, *args, **kwargs):
    return self.make_batchs(*args, **kwargs)

言ってしまえば、訓練データムを格納するリストの構造を変えているだけである。

4-6-7. make_batch_tensor メソッド

make_batch_tensor メソッド は バッチ:

バッチ
list[              # バッチ
    tuple[         # 訓練データム
        list[str], # 問題
        str        # 正解
    ]
]

を受け取り、単語IDのテンソルのタプル:

テンソルのタプル
tuple[
    torch.Tensor(                        # 問題
        shape=(batch_size, content_len),
        dtype=torch.long
    ), 
    torch.Tensor(                        # 正解
        shape=(batch_size,),
        dtype=torch.long
    )              
]

に変換して返す。
オプション引数strange や 学習器のデータ属性 learner.strange によっては、
未知語を検知した場合にそれらをモデルの語彙に追加する。

フック on_before_word_added, on_word_added, on_strange_ignoredを発火し得る。


まず、引数に異常がないか確認しよう。

batch が空なら ValueError を投げる。
未知語に対するポリシー strange が 4-6-1節で紹介した "add", "ignore", "error" のいずれか、あるいは None でないなら ValueError を投げる。
strange"add", "ignore", "error" のいずれかなら、当該ポリシーを採用し、
None なら、データ属性 self.strange のものを採用する。

make_batch_tensorメソッドで引数に異常がないか確認する
def make_batch_tensor(self, batch, strange=None):
    batch_size = len(batch)
    if batch_size == 0:
        raise ValueError("empty batch")
    
    if isinstance(batch, tuple) and len(batch) == 2:
        batch = [batch]
        
    strange = self.strange if strange is None else strange.lower() 
    if strange not in {"add", "ignore", "error"}:
        raise ValueError('strange must be "add", "ignore", or "error"')
    ...

次に、未知語を抽出しよう。

make_batch_tensorメソッドで未知語を抽出する
def make_batch_tensor(self, batch, strange=None):
    ...
    # ここから追加
    unknowns = []
    seen = set()
    for (ctx_words, target_word) in batch:
        for w in list(ctx_words) + [target_word]:
            if w not in self.model.stoi and w not in seen:
                seen.add(w)
                unknowns.append(w)
    ...                

未知語が unknowns に記録される。


次に、未知語を処理しよう。

make_batch_tensorメソッドで未知語を処理する
def make_batch_tensor(self, batch, strange=None):
    ...
    # ここから追加
    if unknowns:
        if strange == "add":
            self.add_words(
                unknowns, rebuild_optimizer=True, update_cache=True
            )
        if strange == "ignore":
            self.on_strange_ignored(unknowns)
            unknowns = set(unknowns) # 高速化
            batch = [
                (
                    [
                        NULL_TOKEN if w in unknowns else w
                        for w in ctx_words
                    ],
                    NULL_TOKEN if target_word in unknowns else target_word
                ) for (ctx_words, target_word) in batch
            ]
        if strange == "error":
            raise ValueError(f"unknown words: {unknowns}")
    ...

未知語の処理方針が "add" の場合、 add_words メソッドを呼ぶ。
ここで on_before_word_added, on_word_added フックが発火する。

未知語の処理方針が "ignore" の場合、 on_strange_ignored フックが発火する。
そしてバッチ内の未知語が null トークンに置き換えられる。

未知語の処理方針が "error" の場合、 ValueError を投げる。


最後に、未知語処理済みのバッチを、テンソルのタプルに変換しよう。

make_batch_tensorメソッドでバッチをテンソルのタプルに変換する
def make_batch_tensor(self, batch, strange=None):
    ...
    # ここから追加
    max_len = max(len(ctx_words) for ctx_words, _ in batch)
    device = self.model.z0.device

    x = torch.full(
        (batch_size, max_len),
        fill_value=NULL_ID,
        dtype=torch.long,
        device=device,
    )
    y = torch.empty(batch_size, dtype=torch.long, device=device)

    for i, (ctx_words, target_word) in enumerate(batch):
        ctx_ids = [self.model.stoi[w] for w in ctx_words]
        if ctx_ids:
            x[i, :len(ctx_ids)] = torch.tensor(
                ctx_ids, dtype=torch.long, device=device
            )
        y[i] = self.model.stoi[target_word]

    return (x, y)

x は問題のテンソル、 y は正解のテンソルである。

x[i, len(ctx_ids):] には、 null ID が詰められる。


make_batch_tensor メソッド全体のコードは次のとおり。

make_batch_tensorメソッド全体のコードを表示
ComplexTaylorWordFunctionLearnerクラスのmake_batch_tensorメソッド
def make_batch_tensor(self, batch, strange=None):
    # 1. 引数に異常がないか確認する
    batch_size = len(batch)
    if batch_size == 0:
        raise ValueError("empty batch")
    
    if isinstance(batch, tuple) and len(batch) == 2:
        batch = [batch]
        
    strange = self.strange if strange is None else strange.lower() 
    if strange not in {"add", "ignore", "error"}:
        raise ValueError('strange must be "add", "ignore", or "error"')

    # 2. 未知語を抽出する
    unknowns = []
    seen = set()
    for (ctx_words, target_word) in batch:
        for w in list(ctx_words) + [target_word]:
            if w not in self.model.stoi and w not in seen:
                seen.add(w)
                unknowns.append(w)
                
    # 3. 未知語を処理する
    if unknowns:
        if strange == "add":
            self.add_words(
                unknowns, rebuild_optimizer=True, update_cache=True
            )
        if strange == "ignore":
            self.on_strange_ignored(unknowns)
            unknowns = set(unknowns) # 高速化
            batch = [
                (
                    [
                        NULL_TOKEN if w in unknowns else w
                        for w in ctx_words
                    ],
                    NULL_TOKEN if target_word in unknowns else target_word
                ) for (ctx_words, target_word) in batch
            ]
        if strange == "error":
            raise ValueError(f"unknown words: {unknowns}")
        
    # 4. バッチをテンソルのタプルに変換する
    max_len = max(len(ctx_words) for ctx_words, _ in batch)
    device = self.model.z0.device

    x = torch.full(
        (batch_size, max_len),
        fill_value=NULL_ID,
        dtype=torch.long,
        device=device,
    )
    y = torch.empty(batch_size, dtype=torch.long, device=device)

    for (i, (ctx_words, target_word)) in enumerate(batch):
        ctx_ids = [self.model.stoi[w] for w in ctx_words]
        if ctx_ids:
            x[i, :len(ctx_ids)] = torch.tensor(
                ctx_ids, dtype=torch.long, device=device
            )
        y[i] = self.model.stoi[target_word]

    return (x, y)

( make_batch_tensor メソッド全体のコードを表示 ここまで )

4-6-8. update_cache メソッド

update_cache メソッドは、モデルのキャッシュを更新する。何も受け取らず、何も返さない。


コードは次のとおり。

ComplexTaylorWordFunctionLearnerクラスのupdate_cacheメソッド
def update_cache(self):
    self.model.update_povm_cache(eps=1e-8)

4-6-9. operator_l2_regularization メソッド

operator_l2_regularization メソッドは、何も受け取らず、
L2正則化されるための、パラメータの絶対値の二乗平均を計算して返す。

(但しオプション引数として、L2正則化のモードを切り替えることもできる。
この場合は4-6-1節における norm_reg_mode に対応する str型を渡す必要がある。
None を受け取った場合(デフォルト)、データ属性 self.norm_reg_mode が使われる)


コードは次のとおり。

ComplexTaylorWordFunctionLearnerクラスのoperator_l2_regularizationメソッド
def operator_l2_regularization(self, mode=None):
    mode = mode or self.norm_reg_mode
    mode = str(mode).lower()
    if mode not in {"a", "next", "both"}:
        raise ValueError("Illegal mode.")
    
    reg = 0.0
    for n in range(self.model.k + 1):
        if   mode == "a":
            params = self.model.A[n]
        elif mode == "next":
            params = self.model.Next[n]
        elif mode == "both":
            params = torch.cat([
                self.model.A[n].flatten(),
                self.model.Next[n].flatten()
            ])
                
        reg = reg + torch.abs(params).pow(2).mean()
    return reg

テイラー次数毎に平均をとって、その和を取っている。

4-6-10. normalize_batch メソッド

normalize_batch メソッドは、様々な形式のバッチ:

入力形式
  tuple[list[str],str]
| tuple[list[int],int]
| list[tuple[list[str],str]]
| list[tuple[list[int],int]]
| tuple[torch.Tensor,torch.Tensor] # 単語IDのテンソルである必要がある

を受け取り、所定のフォーマット:

出力形式
list[tuple[list[str], str]]

に変換したものを返す。


コードは次のとおり。

normalize_batchメソッドのコードを表示
ComplexTaylorWordFunctionLearnerクラスのnormalize_batchメソッド
def normalize_batch(self, batch):
    # -------------------------
    # Tensor batch
    # -------------------------
    if (
        isinstance(batch, tuple)
        and len(batch) == 2
        and isinstance(batch[0], torch.Tensor)
        and isinstance(batch[1], torch.Tensor)
    ):
        x, y = batch

        result = []

        for ctx_ids, target_id in zip(x.tolist(), y.tolist()):
            ctx_words = [
                self.model.itos[i]
                for i in ctx_ids
            ]

            target_word = self.model.itos[target_id]

            result.append((ctx_words, target_word))

        return result

    # -------------------------
    # Single example
    # -------------------------
    if (
        isinstance(batch, tuple)
        and len(batch) == 2
    ):
        batch = [batch]

    if not isinstance(batch, list):
        raise TypeError("unsupported batch format")

    if len(batch) == 0:
        raise ValueError("empty batch")

    # -------------------------
    # ID batch
    # -------------------------
    ctx0, target0 = batch[0]

    is_id_batch = (
        (len(ctx0) > 0 and isinstance(ctx0[0], int))
        or isinstance(target0, int)
    )

    if is_id_batch:
        batch = [
            (
                [self.model.itos[i] for i in ctx_words],
                self.model.itos[target_word],
            )
            for ctx_words, target_word in batch
        ]

    return batch

( normalize_batch メソッドを表示 ここまで)

4-6-11. learn メソッド

learn メソッド はバッチを受け取るか、または何も受け取らず、$1$ ステップ学習する。
ステップ結果の情報を返す。

通常は バッチを受け取る。


バッチを受け取った場合:

  1. モデルのパラメータの勾配をリセットして
  2. バッチから勾配を計算して、
  3. $1$ ステップ実行する。

を行う。(通常モード)

バッチを受け取らない場合:

  1. モデルのパラメータの勾配をリセットせず
  2. $1$ ステップ実行する。

を行う。(勾配蓄積モード)
こちらは、「勾配蓄積」というテクニックを使う場合に有用である。

勾配蓄積とは、巨大なバッチを表現したいがメモリが足りない場合に、
勾配作成を1つのバッチで行うのではなく、複数回のバッチに分けて行い、
あたかも1つの巨大なバッチであるかのように見せかけて学習するテクニックである。

通常は:

  1. learn(batch0) を呼ぶ
    1. 勾配をリセットする
    2. batch0 から勾配を計算する
    3. 勾配に基づいて $1$ ステップ学習する

とやるが、勾配蓄積をしたい場合は:

  1. zero_grad() を呼ぶ
    1. 勾配をリセットする
  2. for i in range(x): accumulate_grad(batch{i}) を呼ぶ
    1. batch0 から勾配を計算したものを加算する
    2. batch1 から勾配を計算したものを加算する
    3. ...
    4. batch{x-1} から勾配を計算したものを加算する
  3. learn((引数なし)) を呼ぶ
    1. 勾配(batch0batch{x-1})に基づいて $1$ステップ学習する

という手順を踏むことになる。


まずは、オプティマイザが同期されていることを確認しよう。
オプティマイザがセットされていなかったり、モデルの語彙を変更したために古くなっている場合は、
RuntimeError を投げる。

learnメソッドでオプティマイザの同期を確認する
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    if self.opt is None:
        raise RuntimeError(
            "learner.opt is None. Set an optimizer before leran()."
        )

    if not self.optimizer_rebuilt:
        raise RuntimeError(
            "Optimizer is expired. Call set_optimizer()."
        )
    ...

次に、update_cache 引数に異常がないか確認しよう。

その他の引数については、
lr, weight_decay, clip_grad_norm → 確認を省略している。
strange, norm_reg_mode, lambda_normaccumulate_grad メソッドに任せている

learnメソッドでupdate_cache引数を確認する
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    if update_cache is None:
        update_cache_before_learning = \
            self.auto_update_cache_before_learning
        update_cache_after_learning = \
            self.auto_update_cache_after_learning
    else:
        err = ValueError("update_cache argument is illegal.")
        if   isinstance(update_cache, bool):
            before = bool(update_cache)
            after  = bool(update_cache)
        elif isinstance(update_cache, str):
            update_cache = str(update_cache).lower()
            if update_cache not in {"before", "after", "both", "none"}:
                raise err
            before = (update_cache in {"both", "before"})
            after  = (update_cache in {"both", "after"})
        else:
            raise err

        update_cache_before_learning = before
        update_cache_after_learning = after
    ...

update_cache 引数は、4-6-1節(コンストラクタ)の auto_update_cache 引数と同じ仕様であるか、
None(デフォルト) である必要がある。

update_cache 引数に応じて、学習前後にモデルのキャッシュを更新するかどうかが決まる。
但し None の場合は、データ属性を元に決まる。

そしてキャッシュを更新するかの情報は、
update_cache_before_learning 変数と update_cache_after_learning 変数に記録される。


次に、 update_cache_before_learningTrue なら
モデルのキャッシュを更新しよう。

learnメソッドでモデルのキャッシュを更新する
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    if update_cache_before_learning:
        self.update_cache()
    ...

次に、通常モード(batch あり)と勾配蓄積モード(batch=None)で分岐処理をしよう。
通常モード では勾配をリセットしてから、batchに基づいて勾配を得る。
勾配蓄積モードではなにもしない。

learnメソッドで勾配を扱う
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    if batch is None:
        info = None
    else:
        self.zero_grad()
        info = self.accumulate_grad(
            batch, strange=strange,
            norm_reg_mode=norm_reg_mode, lambda_norm=lambda_norm
        )
    ...

通常モードでは accumulate_grad メソッドを呼ぶことで勾配を得ている。
勾配を得るタイミングで、損失関数の値などの情報が得られるので、それを info に入れている。

勾配蓄積モードでは info = None としている。


次に、clip_grad_norm 引数によっては、
勾配を縮小しよう。

clip_grad_normNone の場合は、self.clip_grad_norm が代わりに使われる。
self.clip_grad_norm および Noneでない場合の clip_grad_norm 引数 は 4-6-1節(コンストラクタ) の clip_grad_norm 引数のとおりの仕様とする。

learnメソッドで勾配を縮小する
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    clip_grad_norm = \
        self.clip_grad_norm if clip_grad_norm is None else clip_grad_norm
    if isinstance(clip_grad_norm, bool) and clip_grad_norm:
        clip_grad_norm = 1.0
    if not isinstance(clip_grad_norm, bool):
        torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), clip_grad_norm
        )
    ...

次に 勾配に基づいて1ステップ学習(モデルのパラメータを更新)しよう。

learnメソッドで1ステップ進める
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    if lr is None and weight_decay is None:
        self.opt.step()
    else:
        old_wds = [group["weight_decay"] for group in self.opt.param_groups]
        old_lrs = [group["lr"] for group in self.opt.param_groups]

        try:
            for (i,group) in enumerate(self.opt.param_groups):
                group["weight_decay"] = float(
                    weight_decay(group) if callable(weight_decay) else (
                        old_wds[i] if weight_decay is None else weight_decay
                    )
                )
                group["lr"] = float(
                    lr(group) if callable(lr) else (
                        old_lrs[i] if lr is None else lr
                    )
                )
            self.opt.step()
        finally:
            for (group, old_lr, old_wd) in zip(
                self.opt.param_groups, old_lrs, old_wds
            ):
                group["weight_decay"] = old_wd
                group["lr"] = old_lr
                
    self.model.reset_null()
    ...

基本的には lr=None, weight_decay=None なので、単に self.opt.step() してから
self.model.reset_null() しているだけだ。

但し、「今回だけ学習率または重み減衰を変えたい」という要請に応えるべく、
lrweight_decay を引数でいじれる仕様にしている。

また、self.opt.step()したら、その直後に self.model.reset_null() することを忘れてはならない。


次に、update_cache_after_learningTrue なら、
モデルのキャッシュを更新しよう。

learnメソッドでモデルのキャッシュを更新する
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    if update_cache_after_learning:
        self.update_cache()
    ...

最後に、ステップ結果の情報を返そう。

learnメソッドでステップ結果の情報を返す
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    ...
    # ここから追加
    return info

learn メソッド全体のコードは次のとおり。

learnメソッド全体のコードを表示
ComplexTaylorWordFunctionLearnerクラスのlearnメソッド
def learn(
    self, batch=None, lr=None, weight_decay=None,
    strange=None, norm_reg_mode=None,
    lambda_norm=None, clip_grad_norm=None, update_cache=None
):
    # 1. オプティマイザが同期されているか確認する
    if self.opt is None:
        raise RuntimeError(
            "learner.opt is None. Set an optimizer before leran()."
        )

    if not self.optimizer_rebuilt:
        raise RuntimeError(
            "Optimizer is expired. Call set_optimizer()."
        )

    # 2. update_cache 引数に異常がないか確認する
    if update_cache is None:
        update_cache_before_learning = \
            self.auto_update_cache_before_learning
        update_cache_after_learning = \
            self.auto_update_cache_after_learning
    else:
        err = ValueError("update_cache argument is illegal.")
        if   isinstance(update_cache, bool):
            before = bool(update_cache)
            after  = bool(update_cache)
        elif isinstance(update_cache, str):
            update_cache = str(update_cache).lower()
            if update_cache not in {"before", "after", "both", "none"}:
                raise err
            before = (update_cache in {"both", "before"})
            after  = (update_cache in {"both", "after"})
        else:
            raise err

        update_cache_before_learning = before
        update_cache_after_learning = after

    # 3. 必要ならモデルのキャッシュを更新する
    if update_cache_before_learning:
        self.update_cache()

    # 4. 通常モードなら:
    #        勾配をリセット→batchによる勾配を得る
    #    勾配蓄積モードなら:
    #        何もしない
    if batch is None:
        info = None
    else:
        self.zero_grad()
        info = self.accumulate_grad(
            batch, strange=strange,
            norm_reg_mode=norm_reg_mode, lambda_norm=lambda_norm
        )

    # 5. 必要なら勾配を縮小する
    clip_grad_norm = \
        self.clip_grad_norm if clip_grad_norm is None else clip_grad_norm
    if isinstance(clip_grad_norm, bool) and clip_grad_norm:
        clip_grad_norm = 1.0
    if not isinstance(clip_grad_norm, bool):
        torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), clip_grad_norm
        )

    # 6. 勾配に基づいて1ステップ重みを更新する
    if lr is None and weight_decay is None:
        self.opt.step()
    else:
        old_wds = [group["weight_decay"] for group in self.opt.param_groups]
        old_lrs = [group["lr"] for group in self.opt.param_groups]

        try:
            for (i,group) in enumerate(self.opt.param_groups):
                group["weight_decay"] = float(
                    weight_decay(group) if callable(weight_decay) else (
                        old_wds[i] if weight_decay is None else weight_decay
                    )
                )
                group["lr"] = float(
                    lr(group) if callable(lr) else (
                        old_lrs[i] if lr is None else lr
                    )
                )
            self.opt.step()
        finally:
            for (group, old_lr, old_wd) in zip(
                self.opt.param_groups, old_lrs, old_wds
            ):
                group["weight_decay"] = old_wd
                group["lr"] = old_lr
                
    self.model.reset_null()

    # 7. 必要ならモデルのキャッシュを更新する
    if update_cache_after_learning:
        self.update_cache()

    # 8. ステップ結果の情報を返す
    return info

(learn メソッド全体のコードを表示 ここまで)

4-6-12. zero_grad メソッド

zero_grad メソッド は、オプティマイザの勾配をリセットする。
何も受け取らず、何も返さない。


コードは次のとおり。

ComplexTaylorWordFunctionLearnerクラスのzero_gradメソッド
def zero_grad(self):
    self.opt.zero_grad()

4-6-13. accumulate_grad メソッド

accumulate_grad メソッド はバッチを受け取り、勾配を蓄積する。
そこで得られる損失関数の値などの情報を返す。


learn メソッドと異なり、モデルのパラメータの更新 (learner.opt.step()) は行わない。


まずは、learnメソッドでもやったように、オプティマイザが同期されていることを確認しよう。

accumulate_gradメソッドでオプティマイザの同期を確認する
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    if self.opt is None:
        raise RuntimeError(
            "learner.opt is None. "
            "Set an optimizer before accumulate_grad()."
        )

    if not self.optimizer_rebuilt:
        raise RuntimeError(
            "Optimizer is expired. Call set_optimizer()."
        )
    ...

次に、norm_reg_mode 引数に異常がないか確認しよう。

その他の引数については、
strangemake_batch_tensor メソッドに任せている。
lambda_norm → 確認を省略している。

accumulate_gradメソッドでnorm_reg_mode引数を確認する
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    ...
    # ここから追加
    norm_reg_mode = norm_reg_mode or self.norm_reg_mode
    norm_reg_mode = str(norm_reg_mode).lower()
    if norm_reg_mode not in {"a", "next", "both"}:
        raise ValueError("Illegal norm_reg_mode.")
    ...

norm_reg_mode 引数は、4-6-1節(コンストラクタ)の norm_reg_mode 引数と同じ仕様であるか、
None(デフォルト) である必要がある。

norm_reg_mode 引数に応じて、L2正則化の際に見るモデルパラメータが決まる。
但し None の場合は、データ属性を元に決まる。


次に、バッチをテンソルのタプルに変換しよう。

accumulate_gradメソッドでバッチをテンソルのタプルに変換する
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    ...
    # ここから追加
    batch = self.normalize_batch(batch)
    (x, y) = self.make_batch_tensor(batch, strange=strange)
    ...

x が問題のテンソル, y が正解のテンソルである。


次に、モデルに問題を入力し、その出力と正解のズレから損失を求めよう。
さらに、L2正則化の項も損失に加えよう。

accumulate_gradメソッドで損失を計算する
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    ...
    # ここから追加
    lambda_norm = self.lambda_norm if lambda_norm is None else lambda_norm 

    log_probs = self.model.forward_sequence(x)
    ce_loss = F.nll_loss(log_probs, y)
    norm_reg = self.operator_l2_regularization(mode=norm_reg_mode)
    loss = ce_loss + lambda_norm * norm_reg
    ...

損失関数には、前回の記事の4-1-2節で述べた通り、 nll_loss を使うことで交差エントロピー損失を採用している。(ce_loss)

ここに、 L2正則化項 lambda_norm + self.operator_l2_regularization(mode=norm_reg_mode) を加えたものを、最終的な損失関数 loss としている。


次に、勾配蓄積をしよう。

accumulate_gradメソッドで勾配蓄積する
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    ...
    # ここから追加
    loss.backward()
    ...

loss.backward() をすることで、各パラメータの勾配情報に:

$$\frac{\partial {\rm loss}}{\partial {パラメータ}}$$

が加算される。

もう一度言うが、上書きではなくて、加算だ。

これが勾配「蓄積」と呼ばれる所以である。


最後に、損失関数の値などの情報を返そう。

accumulate_gradメソッドで情報を返す
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    ...
    # ここから追加
    return {
        "loss": loss.item(),
        "ce_loss": ce_loss.item(),
        "norm_reg": norm_reg.item(),
        "batch_size": int(y.shape[0]),
    }

accumulate_grad メソッド全体のコードは次のとおり。

accumulate_gradメソッド全体のコードを表示
ComplexTaylorWordFunctionLearnerクラスのaccumulate_gradメソッド
def accumulate_grad(
    self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
):
    # 1. オプティマイザが同期されているか確認する
    if self.opt is None:
        raise RuntimeError(
            "learner.opt is None. "
            "Set an optimizer before accumulate_grad()."
        )

    if not self.optimizer_rebuilt:
        raise RuntimeError(
            "Optimizer is expired. Call set_optimizer()."
        )

    # 2. norm_reg_mode 引数に異常がないか確認する
    norm_reg_mode = norm_reg_mode or self.norm_reg_mode
    norm_reg_mode = str(norm_reg_mode).lower()
    if norm_reg_mode not in {"a", "next", "both"}:
        raise ValueError("Illegal norm_reg_mode.")

    # 3. バッチをテンソルのタプルに変換する
    batch = self.normalize_batch(batch)
    (x, y) = self.make_batch_tensor(batch, strange=strange)

    # 4. 損失を計算する
    lambda_norm = self.lambda_norm if lambda_norm is None else lambda_norm 

    log_probs = self.model.forward_sequence(x)
    ce_loss = F.nll_loss(log_probs, y)
    norm_reg = self.operator_l2_regularization(mode=norm_reg_mode)
    loss = ce_loss + lambda_norm * norm_reg

    # 5. 勾配蓄積する
    loss.backward()

    # 6. 損失の値などの情報を返す
    return {
        "loss": loss.item(),
        "ce_loss": ce_loss.item(),
        "norm_reg": norm_reg.item(),
        "batch_size": int(y.shape[0]),
    }

4-7. 学習器クラス全体のコード

長いので、折りたたむ。
このコードを ComplexTaylorWordFunctionLearner.py としてカレントディレクトリに保存する。

コードを表示
ComplexTaylorWordFunctionLearner.py
import math
import random
import torch
import torch.nn as nn
import torch.nn.functional as F

from ComplexTaylorWordFunctionLM import (
    ComplexTaylorWordFunctionLM, NULL_TOKEN, NULL_ID, MAKE_A
)

# =========================
# Learner class
# =========================

class ComplexTaylorWordFunctionLearner:
    """Training helper for ComplexTaylorWordFunctionLM.

    This class owns:

        - vocabulary construction/expansion/removal
        - context-example construction
        - mini-batch construction
        - optimizer setup
        - one-step learning
        - optional POVM-cache refreshing
        - unknown-word handling

    The model computation itself remains in ComplexTaylorWordFunctionLM.

    Attributes:
        model (ComplexTaylorWordFunctionLM): the language model being trained

        strange (str):
            Unknown-word handling policy.

            Valid values are:

                "add"
                    Add unknown words to the vocabulary and model.

                "ignore"
                    Replace unknown words with <Null>.

                "error"
                    Raise ValueError when unknown words are found.
        
        auto_update_cache_before_learning (bool):
            whether learn() automatically refreshes the model cache
            before the optimization step

        auto_update_cache_after_learning (bool):
            whether learn() automatically refreshes the model cache
            after the optimization step

        norm_reg_mode (str):            
            Control which model parameters to be L2 regularized.

            "a"   : A parameter only
            "next": Next parameter only
            "both": both A and Next parameter

        lambda_norm (float): a hyper parameter of L2 regularization

        clip_grad_norm (float|bool, optional):
            A hyper parameter that specifies the threshold
            used to scale down gradients when their norm exceeds it.

            If False, scaling down will not be done.

        optimizer_rebuilt (bool):
            Whether the optimizer is synchronized with the current
            model parameters.

            Vocabulary expansion/removal replaces parameters in the model.
            After such operations, this flag becomes False
            until set_optimizer() is called.

        examples (list[tuple[list[str],str]]):
            Training examples.

            Each element is

                (context_words, target_word)

            where

                context_words : list[str]
                target_word   : str

        opt (torch.optim.Optimizer | None):
            Optimizer used by learn().

            May be replaced freely by the caller.
    """

    def __init__(
        self, sentences, d=16, k=4, next_identity_init=True,
        strange="error", optimizer="AdamW",
        lr=7e-4, weight_decay=0.0, lambda_norm=1.0, clip_grad_norm=1.0,
        context_len=8, auto_update_cache="after", norm_reg_mode="both",
        device=None
    ):
        """Initialize the training helper of
        the Complex Taylor Word Function Language Model.

        Args:
            sentences (list[list[str]]):
                Sentence List.

                A sentence should be represented as a word list.
                A word should be str.

            d (int): dimension of the semantic state space of model

            k (int): maximut Taylor order of model

            next_identity_init (bool, optional):
                Control initialization of virtual word operator of model.

                Default value is True.

            strange (str, optional):
                Unknown-word handling policy of learning.

                Valid values are:

                    "add"
                        Add unknown words to the vocabulary and model.

                    "ignore"
                        Replace unknown words with <Null>.

                    "error"
                        Raise ValueError when unknown words are found.

                Default value is "error".

            optimizer (str| type[torch.optim.Optimizer] |None, optional): 
                Default optimizer class.

                The caller can still freely overwrite learner.opt afterwards.

                If str, "AdamW", "Adam", "SGD" is valid.
                If None, no optimizer will be set.

                Default value is "AdamW".

            lr (float, optional):
                Learning rate of optimizer.

                Default value is 7e-4.

            weight_decay (float, optional):
                weight decay of optimizer.

                Default value is 0.0.

            lambda_norm (float, optional):
                A hyper parameter of L2 regularization.

                Default value is 1.0.

            clip_grad_norm (float|bool, optional):
                A hyper parameter that specifies the threshold
                used to scale down gradients when their norm exceeds it.

                If False, scaling down will not be done.
                If True, interpreted as 1.0.

                Default value is 1.0.

            context_len (int, optional):
                Maximum context length used when constructing
                training examples.

                For each target word, at most the previous
                context_len words are used as context.

                Example:

                    sentence = ["I", "love", "you", "very", "much"]

                    context_len = 4

                generates examples such as
                
                    (["<Null>"], "I")
                    (["<Null>", "I"], "love")
                    (["<Null>", "I", "love"], "you")
                    (["<Null>", "I", "love", "you"], "very")
                    (["I", "love", "you", "very"], "much")
                    (["love", "you", "very", "much"], "<Null>").

                Default value is 8.

            auto_update_cache (bool|str, optional):
                Control whether update cache of model automatically
                every time before and after learn.

                If bool:
                    True:  cache will be updated both before and after learn.
                    False: cache won't be updated automatically

                If str:
                    "before": cache will be updated only before learn.
                    "after" : cache will be updated only after learn.
                    "both"  : cache will be updated both before and after learn.
                    "none"  : cache won't be updated automatically
                    else    : raise ValueError.

                Dafault value is "after".

            norm_reg_mode (str, optional):
            
                Control which model parameters to be L2 regularized.

                "A"   : A parameter only
                "Next": Next parameter only
                "both": both A and Next parameter
                else  : raise ValueError

            Default Value is "both".
                

            device (str|torch.device|None, optional):
                Device on which the model is allocated.
                
                Examples:

                    "cpu"
                    "cuda"
                    "cuda:0"
                    torch.device("cuda")

                If None, the model remains on the default device.

                Default value is None.
                
        Raises:
            ValueError: if strange is illegal
            ValueError: if auto_update_cache is illegal
            ValueError: if norm_reg_mode is illegal
            ValueError:
                if d > V so that frame-normalized POVM couldn't span the space
            Exception: if optimizer initializer raises it
        """
        # 1. 設定, フラグ
        strange = str(strange).lower()
        if strange not in {"add", "ignore", "error"}:
            raise ValueError('strange must be "add", "ignore", or "error"')
        self.strange = strange
        
        err = ValueError("auto_update_cache argument is illegal.")
        if   isinstance(auto_update_cache, bool):
            before = bool(auto_update_cache)
            after  = bool(auto_update_cache)
        elif isinstance(auto_update_cache, str):
            auto_update_cache = str(auto_update_cache).lower()
            if auto_update_cache not in {"before", "after", "both", "none"}:
                raise err
            before = (auto_update_cache in {"both", "before"})
            after  = (auto_update_cache in {"both", "after"})
        else:
            raise err

        self.auto_update_cache_before_learning = before
        self.auto_update_cache_after_learning = after

        norm_reg_mode = str(norm_reg_mode).lower()        
        if norm_reg_mode not in {"a", "next", "both"}:
            raise ValueError("norm_reg_mode argument is illegal.")
        self.norm_reg_mode = norm_reg_mode
        
        self.lambda_norm = lambda_norm
        
        if isinstance(clip_grad_norm, bool) and clip_grad_norm:
            clip_grad_norm = 1.0
        if clip_grad_norm is None:
            clip_grad_norm = False
        self.clip_grad_norm = clip_grad_norm
                
        self.optimizer_rebuilt = False

        # 2. モデル
        words = []
        seen = set()
        for sentence in sentences:
            for word in sentence:
                if word != NULL_TOKEN and word not in seen:
                    seen.add(word)
                    words.append(word)

        vocab = sorted(words)
        vocab.insert(NULL_ID, NULL_TOKEN)

        self.model = ComplexTaylorWordFunctionLM(
            vocab, d=d, k=k, next_identity_init=next_identity_init,
        )

        if device is not None:
            self.model.to(torch.device(device))
        self.model.reset_null()

        # 3. 訓練データ
        self.examples = []
        self.examples += self.make_examples(sentences, context_len)

        # 4. オプティマイザ
        self.opt = None
        self.set_optimizer(
            optimizer=optimizer,
            lr=lr,
            weight_decay=weight_decay,
        )


    # -------------------------
    # Optimizer
    # -------------------------

    def set_optimizer(
        self, optimizer="AdamW", lr=7e-4, weight_decay=0.0, **kwargs
    ):
        """Set a default optimizer.

        The caller can still freely overwrite learner.opt afterwards.

        Args:
            optimizer (str| type[torch.optim.Optimizer] |None, optional): 
                Optimizer class.

                If str, "AdamW", "Adam", "SGD" is valid.
                If None, optimizer will be removed from learner.

                Default value is "AdamW".

            lr (float, optional):
                Learning rate of optimizer.

                Default value is 7e-4.

            weight_decay (float, optional):
                weight decay of optimizer.

                Default value is 0.0.

            **kwargs: will be all passed to optimizer

        Returns:
            torch.optim.Optimizer|None:
                The optimizer assigned to this learner.
                If None, optimizer has been removed from learner.

        Raises:
            ValueError: if optimizer str is invalid
            Exception: if optimizer initializer raises it
        """
        if optimizer is None:
            self.opt = None
            return None

        if isinstance(optimizer, str):
            name = optimizer.lower()
            if name == "adamw":
                opt_cls = torch.optim.AdamW
            elif name == "adam":
                opt_cls = torch.optim.Adam
            elif name == "sgd":
                opt_cls = torch.optim.SGD
            else:
                raise ValueError(f"unknown optimizer: {optimizer}")
        else:
            opt_cls = optimizer

        self.opt = opt_cls(
            [p for p in self.model.parameters() if p.requires_grad],
            lr=lr,
            weight_decay=weight_decay,
            **kwargs,
        )
        self.optimizer_rebuilt = True
        return self.opt

    # -------------------------
    # Unknown words / vocabulary expansion
    # -------------------------

    def add_words(self, words, rebuild_optimizer=True, update_cache=True):
        """Append words to the vocabulary and expand the model.

        Calls:
            self.on_before_word_added(added_words)
            self.on_word_added(added_word_and_ids)

        Args:
            words (list[str] | str): words to be added
            
            rebuild_optimizer (bool, optional):
                Whether to rebuild optimizer automatically after words added.

                Adding words replaces nn.Parameter objects in self.model.A,
                so an existing optimizer would otherwise keep references
                to old parameters.

                If False, you must call set_optimizer() manually
                after using this method and before calling learn()
                or accumulate_grad(), or you will receive RuntimeError.

                Default value is True.

            update_cache (bool, optional):
                Whether to update model cache automatically after words added.

                Default value is True.

        Returns:
            list[str]: words actually added to the vocabulary

        Raises:
            Exception: if hooks raise it
            Exception:
                if rebuild_optimizer is True and optimizer initalizer raises it
        """
        # 1. フラグを立てる
        self.optimizer_rebuilt = False

        # 2. 未知語を抽出する
        if isinstance(words, str):
            words = [words]
        
        stoi = dict(self.model.stoi)
        itos = dict(self.model.itos)
        vocab = list(itos.values())
        new_words = []
        for w in words:
            if w not in stoi:
                stoi[w] = len(vocab)
                itos[len(vocab)] = w
                vocab.append(w)
                new_words.append(w)

        # 3. フック on_before_word_added を発火する
        self.on_before_word_added(new_words)

        # 4. モデルに未知語を追加する
        old_V = self.model.V
        add_V = len(new_words)
        new_V = old_V + add_V

        self.model.V = new_V
        self.model.stoi = dict(stoi)
        self.model.itos = dict(itos)

        for n in range(self.model.k + 1):
            old_A = self.model.A[n].detach()
            new_A_rows = MAKE_A(add_V, self.model.d, n).to(
                dtype=old_A.dtype, device=old_A.device
            )
            expanded_A = torch.cat([old_A, new_A_rows], dim=0)
            self.model.A[n] = nn.Parameter(expanded_A)

        old_meanings = self.model.word_meanings_povm.detach()
        extra_meanings = torch.zeros(
            add_V,
            self.model.d,
            dtype=old_meanings.dtype,
            device=old_meanings.device,
        )
        self.model.word_meanings_povm = torch.cat(
            [old_meanings, extra_meanings], dim=0
        )

        self.model.reset_null()

        # 5. 自動処理をする
        if rebuild_optimizer:
            self.set_optimizer(
                optimizer=self.opt.__class__,
                lr=self.opt.defaults['lr'],
                weight_decay=self.opt.defaults['weight_decay'],
            )
        if update_cache:
            self.update_cache()

        # 6. フック on_word_added を発火する
        new_ids = [stoi[w] for w in new_words]
        self.on_word_added(list(zip(new_words, new_ids)))

        # 7. 未知語のリストを返す
        return new_words


    def remove_words(
        self, words, example_policy="drop",
        rebuild_optimizer=True, update_cache=True
    ):
        """Remove words from the vocabulary and shrink the model.

        Calls:
            self.on_before_word_removed(removed_words)
            self.on_word_removed(removed_words)

        Args:
            words (list[str] | str): words to be removed

            example_policy (str, optional):
                How to update self.examples after words are removed.

                - "drop":
                    Remove examples whose context or target contains
                    removed words.

                - "null":
                    Replace removed words in examples with <Null>.

                - "keep":
                    Keep examples unchanged.
                    Be careful: learn(..., strange="add") may add the
                    removed words again.

                Default value is "drop".
                
            rebuild_optimizer (bool, optional):
                Whether to rebuild optimizer automatically after words removed.

                Removing words replaces nn.Parameter objects in self.model.A,
                so an existing optimizer would otherwise keep references
                to old parameters.

                If False, you must call set_optimizer() manually
                after using this method and before calling learn()
                or accumulate_grad(), or you will receive RuntimeError.

                Default value is True.

            update_cache (bool, optional):
                Whether to update model cache automatically after words removed.

                Default value is True.

        Returns:
            list[str]: words actually removed from the vocabulary
            
        Raises:
            ValueError: if illegal example_policy
            ValueError: if <Null> is tried to be removed
            Exception: if hooks raise it
            RuntimeError:
                if a word cannot be deleted because doing would result in d > v
            Exception:
                if rebuild_optimizer is True and optimizer initalizer raises it
        """
        # 1. フラグを立てる
        self.optimizer_rebuilt = False

        # 2. 引数に異常がないか確認する
        if isinstance(words, str):
            words = [words]

        example_policy = example_policy.lower()
        if example_policy not in {"drop", "null", "keep"}:
            raise ValueError('example_policy must be "drop", "null", or "keep"')

        removes = set(words)

        if NULL_TOKEN in removes:
            raise ValueError(f"{NULL_TOKEN} cannot be removed.")

        # 3. 不要語を抽出する
        old_stoi = dict(self.model.stoi)
        old_itos = dict(self.model.itos)

        removed_words = [w for w in old_itos.values() if w in removes]

        # 4. 語彙が少なくなりすぎないか確認する
        new_vocab = [
            old_itos[i]
            for i in range(self.model.V)
            if old_itos[i] not in removes
        ]

        if len(new_vocab) < self.model.d:
            raise RuntimeError(
                "Cannot remove words because the frame-normalized POVM "
                "assumes d <= V. "
                f"After removal, d={self.model.d}, V={len(new_vocab)}."
            )

        # 5. フック on_before_word_removed を発火する
        self.on_before_word_removed(removed_words)

        # 6. モデルから不要語を削除する
        old_to_new = {}
        new_itos = {}
        new_stoi = {}

        for new_id, w in enumerate(new_vocab):
            old_id = old_stoi[w]
            old_to_new[old_id] = new_id
            new_itos[new_id] = w
            new_stoi[w] = new_id

        keep_old_ids = torch.tensor(
            [old_stoi[w] for w in new_vocab],
            dtype=torch.long,
            device=self.model.z0.device,
        )

        self.model.V = len(new_vocab)
        self.model.itos = new_itos
        self.model.stoi = new_stoi

        for n in range(self.model.k + 1):
            old_A = self.model.A[n].detach()
            new_A = old_A.index_select(0, keep_old_ids)
            self.model.A[n] = nn.Parameter(new_A)

        old_meanings = self.model.word_meanings_povm.detach()
        self.model.word_meanings_povm = old_meanings.index_select(
            0, keep_old_ids
        )

        self.model.reset_null()
        
        # 7. 訓練データ内の不要語を処理する
        removeds = set(removed_words)

        if example_policy == "drop":
            self.examples = [
                (ctx_words, target_word)
                for (ctx_words, target_word) in self.examples
                if (
                    target_word not in removeds
                    and all(w not in removeds for w in ctx_words)
                )
            ]

        elif example_policy == "null":
            self.examples = [
                (
                    [
                        NULL_TOKEN if w in removeds else w
                        for w in ctx_words
                    ],
                    NULL_TOKEN if target_word in removeds else target_word,
                )
                for (ctx_words, target_word) in self.examples
            ]

        elif example_policy == "keep":
            pass

        # 8. 自動処理をする
        if rebuild_optimizer:
            self.set_optimizer(
                optimizer=self.opt.__class__,
                lr=self.opt.defaults['lr'],
                weight_decay=self.opt.defaults['weight_decay'],
            )
        if update_cache:
            self.update_cache()

        # 9. フック on_word_removed を発火する
        self.on_word_removed(removed_words)

        # 10. 不要語のリストを返す
        return removed_words

    # -------------------------
    # Examples / Batch construction
    # -------------------------

    def make_examples(self, sentences, context_len):
        """Make examples from sentences but self.examples won't be updated.

        Just calling make_examples() does not affect the batchs
        created by make_batchs().

        To include new sentences in the batchs that will be created by
        make_batchs(), for example:

            learner.examples += learner.make_examples(sentences, 8)

        need to be done.

        Args:
            sentenses (list[list[str]] | list[str]):
                Sentence List or sentence.

                A sentence should be represented as a word list.
                A word should be str.

            context_len (int):
                Maximum context length used when constructing
                training examples.

                For each target word, at most the previous
                context_len words are used as context.

                Example:

                    sentence = ["I", "love", "you", "very", "much"]

                    context_len = 4

                generates examples such as
                
                    (["<Null>"], "I")
                    (["<Null>", "I"], "love")
                    (["<Null>", "I", "love"], "you")
                    (["<Null>", "I", "love", "you"], "very")
                    (["I", "love", "you", "very"], "much")
                    (["love", "you", "very", "much"], "<Null>").

        Returns:
            list[tuple[list[str],str]]: examples
        """
        if sentences and isinstance(sentences[0], str):
            sentences = [sentences]
            
        sentences_without_null = [
            [w for w in s if w != NULL_TOKEN] for s in sentences
        ]
        
        sentences_with_null = [
            [NULL_TOKEN] + s + [NULL_TOKEN] for s in sentences_without_null
        ]

        examples = []

        for sentence in sentences_with_null:
            for t in range(1, len(sentence)):
                start = max(0, t - context_len)
                ctx_words = sentence[start:t]
                target_word = sentence[t]
                examples.append((ctx_words, target_word))

        return examples

    def make_batchs(self, batch_size, shuffle=True):
        """Split all examples into mini-batches for one epoch.

        Each example appears exactly once per call.

        Args:
            batch_size (int): batch size
            
            shuffle (bool, optional):
                Whether examples will be shuffled.

                Default value is True.

        Returns:
            list[list[tuple[list[str], str]]]
        """
        indices = list(range(len(self.examples)))
        if shuffle:
            random.shuffle(indices)

        batchs = []
        for start in range(0, len(indices), batch_size):
            batch_indices = indices[start:start + batch_size]
            batchs.append([self.examples[i] for i in batch_indices])
        return batchs

    def make_batches(self, *args, **kwargs):
        """Alias for make_batchs()."""
        return self.make_batchs(*args, **kwargs)

    def make_batch_tensor(self, batch, strange=None):
        """Convert a word batch into padded id tensors.

        This method:

            1. Detects unknown words.
            2. Handles them according to the strange policy.
            3. Converts words into word IDs.
            4. Pads contexts with NULL_ID.
            5. Returns tensorized mini-batch data.

        Calls:
            self.on_before_word_removed(removed_words)
            self.on_word_removed(removed_words)
            self.on_strange_ignored(ignored_words)

        Args:
            batch (list[tuple[list[str], str]]):
                Batch represented by word strings.

                Each element must be

                    (context_words, target_word)

                where

                    context_words : list[str]
                    target_word  : str

            strange (str|None, optional):
                Unknown-word handling policy.

                Valid str values are:

                    "add"
                        Add unknown words to the vocabulary and model.

                    "ignore"
                        Replace unknown words with <Null>.

                    "error"
                        Raise ValueError when unknown words are found.

                If None, learner.strange will be adopted.

                Default value is None.

        Returns:
            tuple[torch.Tensor, torch.Tensor]: (x,y) where:

                x:
                    Context tensor of shape

                        (batch_size, max_context_length)

                    containing word IDs.

                    Shorter contexts are padded with NULL_ID.

                y:
                    Target tensor of shape

                        (batch_size,)

                    containing target word IDs.

        Raises:
            ValueError: if batch is empty
            ValueError: if strange is not one of "add", "ignore", or "error"
            ValueError: if strange="error" and unknown words are found
        """
        # 1. 引数に異常がないか確認する
        batch_size = len(batch)
        if batch_size == 0:
            raise ValueError("empty batch")
        
        if isinstance(batch, tuple) and len(batch) == 2:
            batch = [batch]
            
        strange = self.strange if strange is None else strange.lower() 
        if strange not in {"add", "ignore", "error"}:
            raise ValueError('strange must be "add", "ignore", or "error"')

        # 2. 未知語を抽出する
        unknowns = []
        seen = set()
        for (ctx_words, target_word) in batch:
            for w in list(ctx_words) + [target_word]:
                if w not in self.model.stoi and w not in seen:
                    seen.add(w)
                    unknowns.append(w)
                    
        # 3. 未知語を処理する
        if unknowns:
            if strange == "add":
                self.add_words(
                    unknowns, rebuild_optimizer=True, update_cache=True
                )
            if strange == "ignore":
                self.on_strange_ignored(unknowns)
                unknowns = set(unknowns) # 高速化
                batch = [
                    (
                        [
                            NULL_TOKEN if w in unknowns else w
                            for w in ctx_words
                        ],
                        NULL_TOKEN if target_word in unknowns else target_word
                    ) for (ctx_words, target_word) in batch
                ]
            if strange == "error":
                raise ValueError(f"unknown words: {unknowns}")
            
        # 4. バッチをテンソルのタプルに変換する
        max_len = max(len(ctx_words) for ctx_words, _ in batch)
        device = self.model.z0.device

        x = torch.full(
            (batch_size, max_len),
            fill_value=NULL_ID,
            dtype=torch.long,
            device=device,
        )
        y = torch.empty(batch_size, dtype=torch.long, device=device)

        for (i, (ctx_words, target_word)) in enumerate(batch):
            ctx_ids = [self.model.stoi[w] for w in ctx_words]
            if ctx_ids:
                x[i, :len(ctx_ids)] = torch.tensor(
                    ctx_ids, dtype=torch.long, device=device
                )
            y[i] = self.model.stoi[target_word]

        return (x, y)

    # -------------------------
    # Learning
    # -------------------------

    def update_cache(self):
        """Update cache of model."""
        self.model.update_povm_cache(eps=1e-8)

    def operator_l2_regularization(self, mode=None):
        """Calculate mean square abs of model's operator's parameter.

        Args:
            mode (str|None, optional):
                Control which model parameters to be calculated.

                If None, learner.norm_reg_mode attribute will be adopted.

                If str:

                    "A"   : A parameter only
                    "Next": Next parameter only
                    "both": both A and Next parameter
                    else  : raise ValueError.

                Default Value is None.

        Returns:
            torch.Tensor:
                Mean square abs of model's operator's parameter.

                Shape is ().

                Dtype is torch.float.

        Raises:
            ValueError: if mode is illegal
        """
        mode = mode or self.norm_reg_mode
        mode = str(mode).lower()
        if mode not in {"a", "next", "both"}:
            raise ValueError("Illegal mode.")
        
        reg = 0.0
        for n in range(self.model.k + 1):
            if   mode == "a":
                params = self.model.A[n]
            elif mode == "next":
                params = self.model.Next[n]
            elif mode == "both":
                params = torch.cat([
                    self.model.A[n].flatten(),
                    self.model.Next[n].flatten()
                ])
                    
            reg = reg + torch.abs(params).pow(2).mean()
        return reg



    def normalize_batch(self, batch):
        """Normalize various batch formats into list[tuple[list[str], str]].

        Tensor inputs are converted back into words using the current
        vocabulary (self.model.itos).

        This method performs only format normalization.
        Unknown-word handling is still delegated to make_batch_tensor().

        Args:
            batch (
                  tuple[list[str],str]
                | tuple[list[int],int]
                | list[tuple[list[str],str]]
                | list[tuple[list[int],int]]
                | tuple[torch.Tensor,torch.Tensor]
            ):
                Supported formats are:
                    
                    Single example (word form): (["I", "love"], "you")
                    Single example (id form):   ([1, 5], 8)
                    
                    Batch (word form):
                    
                        [
                            (["I"], "love"),
                            (["you"], "too"),
                        ]
                        
                    Batch (id form):
                    
                        [
                            ([1], 5),
                            ([2], 7),
                        ]

                    Tensorized batch:

                        (x, y)

                        where

                            x.shape == (batch_size, seq_length)
                            y.shape == (batch_size,)

        Returns:
            list[tuple[list[str], str]]:
                batch converted into a uniform word-based representation

        Raises:
            TypeError: if unsupported batch format
            ValueError: if batch is empty
        """

        # -------------------------
        # Tensor batch
        # -------------------------
        if (
            isinstance(batch, tuple)
            and len(batch) == 2
            and isinstance(batch[0], torch.Tensor)
            and isinstance(batch[1], torch.Tensor)
        ):
            x, y = batch

            result = []

            for ctx_ids, target_id in zip(x.tolist(), y.tolist()):
                ctx_words = [
                    self.model.itos[i]
                    for i in ctx_ids
                ]

                target_word = self.model.itos[target_id]

                result.append((ctx_words, target_word))

            return result

        # -------------------------
        # Single example
        # -------------------------
        if (
            isinstance(batch, tuple)
            and len(batch) == 2
        ):
            batch = [batch]

        if not isinstance(batch, list):
            raise TypeError("unsupported batch format")

        if len(batch) == 0:
            raise ValueError("empty batch")

        # -------------------------
        # ID batch
        # -------------------------
        ctx0, target0 = batch[0]

        is_id_batch = (
            (len(ctx0) > 0 and isinstance(ctx0[0], int))
            or isinstance(target0, int)
        )

        if is_id_batch:
            batch = [
                (
                    [self.model.itos[i] for i in ctx_words],
                    self.model.itos[target_word],
                )
                for ctx_words, target_word in batch
            ]

        return batch

    def learn(
        self, batch=None, lr=None, weight_decay=None,
        strange=None, norm_reg_mode=None,
        lambda_norm=None, clip_grad_norm=None, update_cache=None
    ):
        """Run one optimization step.

        Args:
            batch (
                  tuple[list[str],str]
                | tuple[list[int],int]
                | list[tuple[list[str],str]]
                | list[tuple[list[int],int]]
                | tuple[torch.Tensor,torch.Tensor]
                | None
                , optional
            ):
                Mainly either a list of (context_words, target_word),
                or an already tensorized (x, y) pair.

                Format which normalize_batch() can accept.

                If None, no accumulation of gradients performed,
                and only optimizer step will run without zero_grad.

                If not None, gradients are cleared before accumulation.

                Default value is None.
                

            lr (float|Callable[[dict],float]|None, optional):
                Learning rate.

                If float, will be adopted to all parameters.
                If Callable[[dict],float],
                    input [dict] expects to be parameter group,
                    return float expects to be learning rate.
                If None, settings of learner.opt will be adopted.

                Default value is None.

            weight_decay (float|Callable[[dict],float]|None, optional):
                Weight decay.

                If float, will be adopted to all parameters.
                If Callable[[dict],float],
                    input [dict] expects to be parameter group,
                    return float expects to be weight decay.
                If None, settings of learner.opt will be adopted.

                Default value is None.

            strange (str|None, optional):
                Unknown-word policy.

                    "add"    : append unknown words to the vocabulary/model
                    "ignore" : replace unknown words with <Null>
                    "error"  : raise ValueError when unknown word detected
                    else str : is illegal,
                               so raise ValueError REGARDLESS OF UNKNOWN WORDS
                    None     : learner.strange will be adopted.

                Default value is None.

            norm_reg_mode (str|None, optional):
                Control which model parameters to be L2 regularized.

                If None, learner.norm_reg_mode will be adopted.

                If str:

                    "A"   : A parameter only
                    "Next": Next parameter only
                    "both": both A and Next parameter
                    else  : raise ValueError.

                Default Value is None.
                
            lambda_norm (float, optional):
                A hyper parameter of L2 regularization.

                If None, learner.lambda_norm will be adopted.

                Default value is None.

            clip_grad_norm (float|bool|None, optional):
                A hyper parameter that specifies the threshold
                used to scale down gradients when their norm exceeds it.

                If False, scaling down will not be done.
                If True, interpreted as 1.0.
                
                If None, learner.clip_grad_norm will be adopted.

                Default value is None.

            update_cache (bool|str|None, optional):
                Control whether update cache of model before and after step.

                If None, learner attribute:

                    learner.auto_update_cache_before_learning
                    learner.auto_update_cache_after_learning

                , which may be specified in learner initializer will be adopted.

                If bool:
                    True:  cache will be updated both before and after learn.
                    False: cache won't be updated automatically

                If str:
                    "before": cache will be updated only before learn.
                    "after" : cache will be updated only after learn.
                    "both"  : cache will be updated both before and after learn.
                    "none"  : cache won't be updated automatically
                    else    : raise ValueError.

                Default value is None.

        Returns:
            dict[str, float|int] | None:
                Information dict about the accumulation of gradients
                if batch is not None.

                Keys:
                    "loss" (float):
                        total loss
                            ce_loss + lambda_norm * norm_reg
                    "ce_loss" (float): negative log-likelihood loss
                    "norm_reg" (float): operator L2 regularization term
                    "batch_size" (int):
                        actual number of training examples in the batch

                If batch is None, None.

        Raises:
            RuntimeError:
                If learner.opt is None.

                Call set_optimizer() or assign learner.opt manually
                before calling learn().

            RuntimeError:
                If the optimizer is expired because model parameters were
                replaced after vocabulary expansion/removal.

                Call set_optimizer() before learning.

            ValueError: if norm_reg_mode is illegal
            ValueError: if update_cache is illegal
            
            Exception:
                propagated from normalize_batch(), make_batch_tensor(),
                or hook methods.
        """
        # 1. オプティマイザが同期されているか確認する
        if self.opt is None:
            raise RuntimeError(
                "learner.opt is None. Set an optimizer before leran()."
            )

        if not self.optimizer_rebuilt:
            raise RuntimeError(
                "Optimizer is expired. Call set_optimizer()."
            )

        # 2. update_cache 引数に異常がないか確認する
        if update_cache is None:
            update_cache_before_learning = \
                self.auto_update_cache_before_learning
            update_cache_after_learning = \
                self.auto_update_cache_after_learning
        else:
            err = ValueError("update_cache argument is illegal.")
            if   isinstance(update_cache, bool):
                before = bool(update_cache)
                after  = bool(update_cache)
            elif isinstance(update_cache, str):
                update_cache = str(update_cache).lower()
                if update_cache not in {"before", "after", "both", "none"}:
                    raise err
                before = (update_cache in {"both", "before"})
                after  = (update_cache in {"both", "after"})
            else:
                raise err

            update_cache_before_learning = before
            update_cache_after_learning = after

        # 3. 必要ならモデルのキャッシュを更新する
        if update_cache_before_learning:
            self.update_cache()

        # 4. 通常モードなら:
        #        勾配をリセット→batchによる勾配を得る
        #    勾配蓄積モードなら:
        #        何もしない
        if batch is None:
            info = None
        else:
            self.zero_grad()
            info = self.accumulate_grad(
                batch, strange=strange,
                norm_reg_mode=norm_reg_mode, lambda_norm=lambda_norm
            )

        # 5. 必要なら勾配を縮小する
        clip_grad_norm = \
            self.clip_grad_norm if clip_grad_norm is None else clip_grad_norm
        if isinstance(clip_grad_norm, bool) and clip_grad_norm:
            clip_grad_norm = 1.0
        if not isinstance(clip_grad_norm, bool):
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), clip_grad_norm
            )

        # 6. 勾配に基づいて1ステップ重みを更新する
        if lr is None and weight_decay is None:
            self.opt.step()
        else:
            old_wds = [group["weight_decay"] for group in self.opt.param_groups]
            old_lrs = [group["lr"] for group in self.opt.param_groups]

            try:
                for (i,group) in enumerate(self.opt.param_groups):
                    group["weight_decay"] = float(
                        weight_decay(group) if callable(weight_decay) else (
                            old_wds[i] if weight_decay is None else weight_decay
                        )
                    )
                    group["lr"] = float(
                        lr(group) if callable(lr) else (
                            old_lrs[i] if lr is None else lr
                        )
                    )
                self.opt.step()
            finally:
                for (group, old_lr, old_wd) in zip(
                    self.opt.param_groups, old_lrs, old_wds
                ):
                    group["weight_decay"] = old_wd
                    group["lr"] = old_lr
                    
        self.model.reset_null()

        # 7. 必要ならモデルのキャッシュを更新する
        if update_cache_after_learning:
            self.update_cache()

        # 8. ステップ結果の情報を返す
        return info

    def zero_grad(self):
        """Call self.opt.zero_grad()."""
        self.opt.zero_grad()
        
    def accumulate_grad(
        self, batch, strange=None, norm_reg_mode=None, lambda_norm=None
    ):
        """Run one accumulation of gradients.

        Args:
            batch (
                  tuple[list[str],str]
                | tuple[list[int],int]
                | list[tuple[list[str],str]]
                | list[tuple[list[int],int]]
                | tuple[torch.Tensor,torch.Tensor]
            ):
                Mainly either a list of (context_words, target_word),
                or an already tensorized (x, y) pair.

                Format which normalize_batch() can accept.

            strange (str|None, optional):
                Unknown-word policy.

                    "add"    : append unknown words to the vocabulary/model
                    "ignore" : replace unknown words with <Null>
                    "error"  : raise ValueError when unknown word detected
                    else str : is illegal,
                               so raise ValueError REGARDLESS OF UNKNOWN WORDS
                    None     : learner.strange will be adopted.

                Default value is None.

            norm_reg_mode (str|None, optional):
                Control which model parameters to be L2 regularized.

                If None, learner.norm_reg_mode will be adopted.

                If str:

                    "A"   : A parameter only
                    "Next": Next parameter only
                    "both": both A and Next parameter
                    else  : raise ValueError.

                Default Value is None.
                
            lambda_norm (float, optional):
                A hyper parameter of L2 regularization.

                If None, learner.lambda_norm will be adopted.

                Default value is None.

        Returns:
            dict[str, float|int]:
                Information about the accumulation.

                Keys:
                    "loss" (float):
                        total loss
                            ce_loss + lambda_norm * norm_reg
                    "ce_loss" (float): negative log-likelihood loss
                    "norm_reg" (float): operator L2 regularization term
                    "batch_size" (int):
                        actual number of training examples in the batch

        Raises:
            RuntimeError:
                If learner.opt is None.

                Call set_optimizer() or assign learner.opt manually
                before calling accumulate_grad().

            RuntimeError:
                If the optimizer is expired because model parameters were
                replaced after vocabulary expansion/removal.

                Call set_optimizer() before learning.

            ValueError: if norm_reg_mode is illegal
            
            Exception:
                propagated from normalize_batch(), make_batch_tensor(),
                or hook methods.
        """
        # 1. オプティマイザが同期されているか確認する
        if self.opt is None:
            raise RuntimeError(
                "learner.opt is None. "
                "Set an optimizer before accumulate_grad()."
            )

        if not self.optimizer_rebuilt:
            raise RuntimeError(
                "Optimizer is expired. Call set_optimizer()."
            )

        # 2. norm_reg_mode 引数に異常がないか確認する
        norm_reg_mode = norm_reg_mode or self.norm_reg_mode
        norm_reg_mode = str(norm_reg_mode).lower()
        if norm_reg_mode not in {"a", "next", "both"}:
            raise ValueError("Illegal norm_reg_mode.")

        # 3. バッチをテンソルのタプルに変換する
        batch = self.normalize_batch(batch)
        (x, y) = self.make_batch_tensor(batch, strange=strange)

        # 4. 損失を計算する
        lambda_norm = self.lambda_norm if lambda_norm is None else lambda_norm 

        log_probs = self.model.forward_sequence(x)
        ce_loss = F.nll_loss(log_probs, y)
        norm_reg = self.operator_l2_regularization(mode=norm_reg_mode)
        loss = ce_loss + lambda_norm * norm_reg

        # 5. 勾配蓄積する
        loss.backward()

        # 6. 損失の値などの情報を返す
        return {
            "loss": loss.item(),
            "ce_loss": ce_loss.item(),
            "norm_reg": norm_reg.item(),
            "batch_size": int(y.shape[0]),
        }


    # -------------------------
    # Hooks
    # -------------------------
    def on_before_word_added(self, added_words): pass
    def on_word_added(self, added_word_and_ids): pass
    def on_strange_ignored(self, ignored_words): pass
    def on_before_word_removed(self, removed_words): pass
    def on_word_removed(self, removed_words): pass

(コードを表示 ここまで)

4-8. 学習器クラスのテスト

現状のステータスは次のとおり。

ディレクトリ構成とステータス
(カレントディレクトリ)
├ main.py                               # 😴メインプログラム
├ ComplexTaylorWordFunctionLM.py        # ✅モデルクラス
├ ComplexTaylorWordFunctionLearner.py   # 👀モデルを学習するためのクラス
├ constants.py                          # ✅複数ファイルで共有する定数など (コーパス以外)
├ corpus.py                             # 😴トイコーパス
├ ComplexTaylorWordFunctionEvaluator.py # 😴学習結果や経過を評価するためのクラス
└ test_project.py                       # 🛠️テスト用コード

ステータスの凡例
😴: 未着手
🛠️: 今から作成・編集する
📌: 新たに編集が必要になる
👀: 作成済み(未テスト)
✅: 作成済み (テスト成功済み or テストを実施しない)
⚠️: 問題あり

今回は、 pytest を使って ComplexTaylorWordFunctionLearner.py をテストしていこう。

テストコード test_project.py は次のとおり。

テストコード test_project.py を表示
test_project.py
import pytest
from constants import NULL_TOKEN, NULL_ID
from ComplexTaylorWordFunctionLM import ComplexTaylorWordFunctionLM
from ComplexTaylorWordFunctionLearner import ComplexTaylorWordFunctionLearner
import torch
import numpy as np
import math

# =========================
# Test
# =========================

# utils
def norm(tensor):
    try:
        return torch.linalg.norm(tensor)
    except TypeError:
        return np.linalg.norm(tensor)
normalize = lambda tensor: tensor / norm(tensor)    
assert_close = torch.testing.assert_close
tensor = lambda l, dtype=torch.cfloat: torch.tensor(l, dtype=dtype)
def assert_not_close(a, b, **kwargs):
    with pytest.raises(AssertionError) as exc_info:
        assert_close(a, b, **kwargs)

# Model
@pytest.fixture
def vocab():
    return [NULL_TOKEN, "I", "love", "you", "thank", "very", "much"]

@pytest.fixture
def model(vocab):
    return ComplexTaylorWordFunctionLM(vocab, d=4)

def test_Model_init_error(vocab):
    with pytest.raises(ValueError, match="must have word ID") as excinfo:
        vocab_illegal = list(vocab[1:])
        model = ComplexTaylorWordFunctionLM(vocab_illegal, d=4)

    with pytest.raises(
        ValueError, match="POVM needs enough word meanings"
    ) as excinfo:
        model = ComplexTaylorWordFunctionLM(vocab)

def test_Model_init(vocab, model):
    assert model.V == len(vocab)
    assert all(
        i == i_ and s == s_
        for ((i,s), (s_,i_))
        in zip(model.itos.items(), model.stoi.items())
    )
    assert len(vocab) == len(model.itos)
    assert len(model.itos) == len(model.stoi)
    assert model.d == 4
    assert model.k == 4

    A = model.A
    assert isinstance(A, torch.nn.ParameterList)
    assert len(A) == model.k + 1
    assert isinstance(A[0], torch.nn.Parameter)
    for n in range(model.k + 1):
        assert A[n].shape == (model.V, model.d, model.d ** n)
    print(f"\n{model.A[0]=}\n{model.A[1]=}")

    Next = model.Next
    assert isinstance(Next, torch.nn.ParameterList)
    assert len(Next) == model.k + 1
    assert isinstance(Next[0], torch.nn.Parameter)
    for n in range(model.k + 1):
        assert Next[n].shape == (model.d, model.d ** n)
    print(f"\n{model.Next[0]=}\n{model.Next[1]=}")

    z0 = model.z0
    assert isinstance(z0, torch.Tensor)
    assert z0.shape == (model.d,)
    print(f"\n{model.z0=}")
    
    assert isinstance(model.word_meanings_povm, torch.Tensor)
    assert model.word_meanings_povm.shape == (model.V, model.d)
    assert isinstance(model.G_povm, torch.Tensor)
    assert model.G_povm.shape == (model.d, model.d)
    assert isinstance(model.G_inv_sqrt_povm, torch.Tensor)
    assert model.G_inv_sqrt_povm.shape == (model.d, model.d)

def test_Model_reset_null(model):
    # 1. null を破壊する
    with torch.no_grad():
        model.A[0][NULL_ID,:,:].fill_(3.0+4.0j)

    # 2. null が恒等作用素でなくなったことを確認する
    null_meaning = model.word_meaning([NULL_ID])[0]
    assert torch.linalg.norm(model.z0 - null_meaning) > 1e-8

    # 3. reset_null する
    model.reset_null()

    # 4. null が恒等作用素に戻ったことを確認する
    null_meaning = model.word_meaning([NULL_ID])[0]
    assert norm(model.z0 - null_meaning) < 1e-8

def test_Model_tensor_power(model):
    z0 = [1, 2+3j, -4j, 5j-6]
    z1 = [7, 8, 9, 10]
    z_batch = tensor([z0, z1])

    # 0階冪
    assert_close(model.tensor_power(z_batch, 0), tensor([[1],[1]]))

    # 1階冪
    assert_close(model.tensor_power(z_batch, 1), tensor([z0,z1]))

    # 2階冪
    z0_otimes_z0_flatten = [x * y for x in z0 for y in z0]
    z1_otimes_z1_flatten = [x * y for x in z1 for y in z1]
    assert_close(
        model.tensor_power(z_batch, 2),
        tensor([z0_otimes_z0_flatten, z1_otimes_z1_flatten])
    )

    # 3階冪
    z0_otimes_z0_otimes_z0_flatten = [
        x * y for x in z0_otimes_z0_flatten for y in z0
    ]
    z1_otimes_z1_otimes_z1_flatten = [
        x * y for x in z1_otimes_z1_flatten for y in z1
    ]
    assert_close(
        model.tensor_power(z_batch, 3),
        tensor([z0_otimes_z0_otimes_z0_flatten, z1_otimes_z1_otimes_z1_flatten])
    )

def test_Model_apply_word(model):
    # 正規化されていないz
    z0 = [1, 2+3j, -4j, 5j-6]
    z1 = [7, 8, 9, 10]
    z_batch = tensor([z0, z1])
    
    # 正規化されるため、恒等作用素でも出力が変わるはず
    id_batch = tensor([NULL_ID], int)
    assert_not_close(model.apply_word(id_batch, z_batch), z_batch)

    # zを正規化して再試行
    z0 = normalize(tensor(z0)).tolist()
    z1 = normalize(tensor(z1)).tolist()
    z_batch = tensor([z0, z1])

    # <Null>
    id_batch = tensor([NULL_ID], int)
    assert_close(model.apply_word(id_batch, z_batch), z_batch)

    # <Null>, you
    you_id = model.stoi["you"]
    id_batch = tensor([NULL_ID, you_id], int)
    terms = []
    for n in range(model.k + 1):
        A_you_n = model.A[n][you_id].detach().numpy()
        terms.append(
            np.dot(
                A_you_n, model.tensor_power(z_batch, n)[1].numpy()
            ) / math.factorial(n)
        )
    null_z0 = z0
    you_z1 = normalize(np.sum(terms, axis=0)).tolist()
    y_batch = tensor([null_z0, you_z1])
    
    assert_close(model.apply_word(id_batch, z_batch), y_batch)

def test_Model_sequence_meaning(model):
    I_id     = model.stoi["I"]
    love_id  = model.stoi["love"]
    you_id   = model.stoi["you"]
    thank_id = model.stoi["thank"]
    word_seq_batch = tensor([
        [I_id, love_id, you_id],
        [thank_id, you_id, NULL_ID]
    ], int)

    meaning_batch = model.sequence_meaning(word_seq_batch)
    assert isinstance(meaning_batch, torch.Tensor)
    assert meaning_batch.shape == (2, model.d)
    assert meaning_batch.dtype == torch.cfloat

    meaning_I_love_you = meaning_batch[0]
    z_0_batch = model.z0.unsqueeze(0)
    z_I_batch          = model.apply_word(tensor([I_id],    int), z_0_batch)
    z_I_love_batch     = model.apply_word(tensor([love_id], int), z_I_batch)
    z_I_love_you_batch = model.apply_word(tensor([you_id],  int),
                                                             z_I_love_batch)
    z_I_love_you = z_I_love_you_batch[0]
    assert_close(meaning_I_love_you, z_I_love_you)

def test_Model_apply_next(model):
    z0 = normalize(tensor([1, 2+3j, -4j, 5j-6])).tolist()
    z1 = normalize(tensor([7, 8, 9, 10])).tolist()
    z_batch = tensor([z0, z1])

    assert_close(model.apply_next(z_batch), z_batch)

def test_Model_word_meaning(model):
    id_batch = tensor([1,2], int)
    word_meaning_batch = model.word_meaning(id_batch)

    z0_batch = model.z0.unsqueeze(0)
    word_meaning_batch_by_apply_word = model.apply_word(id_batch, z0_batch)

    assert_close(word_meaning_batch, word_meaning_batch_by_apply_word)

def test_Model_update_povm_cache(model):
    # 1. パラメータをランダムに書き換える
    with torch.no_grad():
        for a in model.A:
            a.normal_(mean=0.0, std=1.0)
    model.reset_null()

    # 2. POVMキャッシュが同期されていないことを確認する
    word_meanings = model.word_meaning(torch.arange(model.V))
    word_meanings_np = word_meanings.detach().numpy()
    G_np = np.dot(word_meanings_np.T, np.conj(word_meanings_np))
    G = torch.from_numpy(G_np)

    assert_not_close(word_meanings, model.word_meanings_povm)
    assert_not_close(G, model.G_povm)
    assert_not_close(
        torch.inverse(G), model.G_inv_sqrt_povm @ model.G_inv_sqrt_povm
    )

    # 3. POVM キャッシュを同期する
    model.update_povm_cache()

    # 4. 同期されたことを確認する
    assert_close(word_meanings, model.word_meanings_povm)
    assert_close(G, model.G_povm)
    assert_close(
        torch.inverse(G), model.G_inv_sqrt_povm @ model.G_inv_sqrt_povm
    )

def assert_probs_batch(probs_batch):
    B = probs_batch.shape[0]
    
    # どの値も0以上か
    assert torch.all(torch.min(probs_batch, dim=1).values >= 0)

    # 和はほとんど1か
    assert_close(
        torch.sum(probs_batch, dim=1), tensor([1] * B),
        check_dtype=False
    )
assert_log_probs_batch = lambda log_pb: assert_probs_batch(torch.exp(log_pb))
    

def test_Model_probs_3_methods(model):
    z0 = normalize(tensor([1, 2+3j, -4j, 5j-6])).tolist()
    z1 = normalize(tensor([7, 8, 9, 10])).tolist()
    z_batch = tensor([z0, z1])
    model.update_povm_cache()

    # 1. born_probs_from_meaning メソッド
    probs_batch = model.born_probs_from_meaning(z_batch)
    assert_probs_batch(probs_batch)

    # 2. amplitudes_from_meaning メソッド
    assert_close(
        torch.abs(model.amplitudes_from_meaning(z_batch))**2, probs_batch,
        check_dtype=False
    )

    # 3. log_probs_from_meaning メソッド
    log_probs_batch = model.log_probs_from_meaning(z_batch)
    assert_close(
        log_probs_batch, torch.log(probs_batch)
    )
    
def test_Model_forward(model):
    model.update_povm_cache()
    
    id_batch = tensor([1,2], int)
    assert_log_probs_batch(model.forward(id_batch))

def test_Model_forward_sequence(model):
    model.update_povm_cache()

    I_id     = model.stoi["I"]
    love_id  = model.stoi["love"]
    you_id   = model.stoi["you"]
    thank_id = model.stoi["thank"]
    word_seq_batch = tensor([
        [I_id, love_id, you_id],
        [thank_id, you_id, NULL_ID]
    ], int)

    log_probs_batch = model.forward_sequence(word_seq_batch)
    assert_log_probs_batch(log_probs_batch)

    log_probs_thank_you_null = log_probs_batch[1]
    z_0_batch = model.z0.unsqueeze(0)
    z_thank_batch          = model.apply_word(tensor([thank_id], int),z_0_batch)
    z_thank_you_batch      = model.apply_word(tensor([you_id],   int),
                                                                  z_thank_batch)
    z_thank_you_null_batch = model.apply_word(tensor([NULL_ID],  int),
                                                              z_thank_you_batch)
    assert_close(
        log_probs_thank_you_null,
        model.log_probs_from_meaning(z_thank_you_null_batch)[0]
    )

def test_Model_povm_normalization_error(model):
    model.update_povm_cache()

    err = model.povm_normalization_error()
    assert isinstance(err, float)
    assert torch.abs(tensor(err)) < 1e-4
    print(f"\n{err=}")

def test_Model_speak(model):
    """
    Only do a quick, basic test since this is not a learning related method.
    """
    model.update_povm_cache()

    assert isinstance(model.speak(8, print_text=False), str)
    for i in range(10):
        print(f"\n({i=},){model.speak(i, print_text=False)=}")

    assert isinstance(model.speak(8, return_list=True, print_text=False), list)
    for i in range(10):
        print(f"\n({i=},){model.speak(i, return_list=True, print_text=False)=}")

    with pytest.raises(ValueError, match="n must be non-negative") as exc_info:
        model.speak(-1, print_text=False)

def test_Model_operator_2_methods(model):
    A_one = model.flat_operator(1, adjust_n=False, normalize=False)
    assert A_one.shape == (model.d * sum(model.d**n for n in range(model.k+1)),)

    A_one_adjusted = model.flat_operator(1, adjust_n=True, normalize=False)
    assert_close(
        A_one / A_one_adjusted,
        tensor([
            (model.d * model.d**n)**0.5
            for n in range(model.k + 1)
            for _ in range(model.d *model.d**n)
        ])
    )

    A_one_normalized = model.flat_operator(1, adjust_n=False, normalize=True)
    assert_close(A_one, A_one_normalized * model.operator_norm(1))

# Learner
@pytest.fixture
def sentences():
    return [
        ["I", "love", "you", "very", "much"],
        ["you", "love", "you"],
        ["thank", "you", "very", "much"]
    ]
@pytest.fixture
def learner(sentences):
    return ComplexTaylorWordFunctionLearner(sentences, d=4)

def test_Learner_init_error(sentences):
    with pytest.raises(ValueError, match="strange must be") as exc_info:
        learner = ComplexTaylorWordFunctionLearner(
            sentences, strange="illegal_strange"
        )

    with pytest.raises(ValueError, match="auto_update_cache") as exc_info:
        learner = ComplexTaylorWordFunctionLearner(
            sentences, auto_update_cache="illegal_auto_update_cache"
        )

    with pytest.raises(ValueError, match="norm_reg_mode") as exc_info:
        learner = ComplexTaylorWordFunctionLearner(
            sentences, norm_reg_mode="illegal_norm_reg_mode"
        )

    with pytest.raises(ValueError, match="POVM needs enough word meanings")\
         as exc_info:
        learner = ComplexTaylorWordFunctionLearner(sentences)

    with pytest.raises(Exception) as exc_info:
        learner = ComplexTaylorWordFunctionLearner(
            sentences, d=4, optimizer="illegal_optimizer"
        )

def test_Learner_init(learner, sentences):
    # strange
    assert learner.strange == "error"

    # auto_update_cache
    assert not learner.auto_update_cache_before_learning
    assert learner.auto_update_cache_after_learning
    
    learner1 = ComplexTaylorWordFunctionLearner(
        sentences, d=4, auto_update_cache=True
    )
    assert learner1.auto_update_cache_before_learning
    assert learner1.auto_update_cache_after_learning

    learner2 = ComplexTaylorWordFunctionLearner(
        sentences, d=4, auto_update_cache=False
    )
    assert not learner2.auto_update_cache_before_learning
    assert not learner2.auto_update_cache_after_learning

    assert learner.model.d == 4
    assert learner.model.k == 4

    assert (['<Null>', 'I', 'love', 'you', 'very', 'much'], '<Null>') \
        in learner.examples
    learner3 = ComplexTaylorWordFunctionLearner(
        sentences, d=4, context_len=4
    )
    assert (['<Null>', 'I', 'love', 'you', 'very', 'much'], '<Null>') \
        not in learner3.examples    
    assert (['love', 'you', 'very', 'much'], '<Null>') \
        in learner3.examples

    assert learner.opt.__class__ == torch.optim.AdamW

def test_Learner_set_optimizer(learner):
    # None
    learner.set_optimizer(None)
    assert learner.opt is None

    # (error)
    with pytest.raises(ValueError, match="unknown optimizer") as exc_info:
        learner.set_optimizer("illegal_optim")

    # "SGD"
    learner.set_optimizer("SGD")
    assert learner.opt.__class__ == torch.optim.SGD

    # Adam
    learner.set_optimizer(torch.optim.Adam)
    assert learner.opt.__class__ == torch.optim.Adam

    # kwargs
    learner.set_optimizer(eps=2.0)
    assert learner.opt.defaults["eps"] == 2.0

def test_Learner_add_words(learner):
    V_old = learner.model.V
    # optimizer_rebuilt flag (1/2)
    assert learner.optimizer_rebuilt
    learner.add_words(["strange0"], rebuild_optimizer=False)
    assert not learner.optimizer_rebuilt
    learner.set_optimizer()
    assert learner.optimizer_rebuilt

    # optimizer_rebuilt flag (2/2)
    result = learner.add_words(["strange1"])
    assert learner.optimizer_rebuilt
    assert result == ["strange1"]

    # hook
    learner.added_words = []
    def hook(added_words):
        learner.added_words += added_words
    learner.on_before_word_added = hook
    learner.add_words(["strange2"])
    assert learner.added_words == ["strange2"]

    # V
    assert learner.model.V == V_old + 3
    result = learner.add_words(["strange2"]) # 既に追加済みの単語
    assert learner.model.V == V_old + 3
    assert result == []

def test_Learner_remove_words(learner):
    # 存在しない単語を削除しても何も起きない
    V_old = learner.model.V
    result = learner.remove_words(["strange0"])
    assert result == []
    assert learner.model.V == V_old

    # example_policy の異常検知
    with pytest.raises(ValueError, match="example_policy") as exc_info:
        learner.remove_words(["I"], example_policy="illegal_value")

    # example_policy = "keep"
    old_examples = list([e for e in learner.examples])
    learner.remove_words(["I"], example_policy="keep")
    assert learner.examples == old_examples
    learner.add_words(["I"])

    # example_policy = "null"
    assert any(["love" in e for e in old_examples])
    result = learner.remove_words(["love"], example_policy="null")
    assert result == ["love"]
    assert not any(["love" in e for e in learner.examples])
    assert len(learner.examples) == len(old_examples)
    assert (['<Null>', 'I', '<Null>', 'you', 'very', 'much'], '<Null>') \
        in learner.examples
    learner.add_words(["love"])
    learner.examples = list([e for e in old_examples])

    # example_policy = "drop"
    assert any(["you" in e for e in old_examples])
    learner.remove_words(["you"])
    assert not any(["you" in e for e in learner.examples])
    assert len(learner.examples) < len(old_examples)
    learner.add_words(["you"])
    learner.examples = list([e for e in old_examples])

    # hook
    learner.removed_words = []
    def hook(removed_words):
        learner.removed_words += removed_words
    learner.on_word_removed = hook
    learner.remove_words(["very", "much"])
    assert set(learner.removed_words) == {"very", "much"}

    # null 削除拒否
    with pytest.raises(ValueError, match="cannot be removed") as exc_info:
        learner.remove_words([NULL_TOKEN])

    # 語彙が小さくなりすぎる
    with pytest.raises(RuntimeError, match="assumes d <= V") as exc_info:
        while learner.model.V:
            id1 = learner.model.itos[1]
            learner.remove_words([id1])

def test_Learner_make_examples(learner):
    sentences = [
        ["I", "love", "you", "very", "much"],
        ["thank", NULL_TOKEN, "you"]
    ]
    examples = learner.make_examples(sentences, context_len=4)
    assert sorted(examples) == sorted([
        ([NULL_TOKEN],                                     "I"),
        ([NULL_TOKEN, "I"],                                "love"),
        ([NULL_TOKEN, "I", "love"],                        "you"),
        ([NULL_TOKEN, "I", "love", "you"],                 "very"),
        ([            "I", "love", "you", "very"],         "much"),
        ([                 "love", "you", "very", "much"], NULL_TOKEN),
        
        ([NULL_TOKEN],                 "thank"),
        ([NULL_TOKEN, "thank"],        "you"),
        ([NULL_TOKEN, "thank", "you"], NULL_TOKEN)
    ]) 

def test_Learner_make_batchs(learner):
    assert max([len(batch) for batch in learner.make_batchs(8)]) == 8

    # シャッフルする場合
    examples = []
    for batch in learner.make_batchs(6):
        examples += batch
    assert sorted(learner.examples) == sorted(examples)
    assert learner.examples != examples

    # シャッフルしない場合
    examples = []
    for batch in learner.make_batchs(6, shuffle=False):
        examples += batch
    assert learner.examples == examples

def test_Learner_make_batch_tensor(learner):
    batch = learner.make_batchs(6)[0]

    # empty batch
    with pytest.raises(ValueError, match="empty batch") as exc_info:
        learner.make_batch_tensor([])

    # strange illegal
    with pytest.raises(ValueError, match="strange") as exc_info:
        learner.make_batch_tensor(batch, strange="illegal_strange")

    # strange error
    with pytest.raises(ValueError, match="unknown words") as exc_info:
        learner.make_batch_tensor([(["未知語0"], "未知語1")], strange="error")

    # strange add
    V_old = learner.model.V
    result = learner.make_batch_tensor(
        [(["未知語0"], "未知語1")] + batch, strange="add"
    )
    assert learner.model.V == V_old + 2
    assert isinstance(result, tuple)
    (問題tensor, 正解tensor) = result
    assert 問題tensor[0][0] == learner.model.stoi["未知語0"]
    assert 正解tensor[0]    == learner.model.stoi["未知語1"]
    assert 問題tensor[1][0] == learner.model.stoi[batch[0][0][0]]
    assert 正解tensor[2]    == learner.model.stoi[batch[1][1]]
    
def test_Learner_update_cache(learner):
    pass

def test_Learner_operator_l2_regularization(learner):
    assert learner.operator_l2_regularization(mode=None) == \
           learner.operator_l2_regularization(mode="both")

    with pytest.raises(ValueError) as exc_info:
        learner.operator_l2_regularization(mode="illegal_mode")

    assert abs(
        learner.operator_l2_regularization(mode="A").item()
        - sum(
            torch.abs(learner.model.A[n]).pow(2).mean()
            for n in range(learner.model.k + 1)
        )
    ) < 1e-8

def test_Learner_normalize_batch(learner):
    batch = learner.make_batchs(6)[0]
    assert learner.normalize_batch(batch) == batch

    stoi = learner.model.stoi
    batch_by_id = [([stoi[s] for s in q], stoi[a]) for (q, a) in batch]
    assert learner.normalize_batch(batch_by_id) == batch

    batch_as_tensor = learner.make_batch_tensor(batch)

    # padding があるため元には戻らない
    assert learner.normalize_batch(batch_as_tensor) != batch

    # padding をちゃんとすれば元に戻る
    q_len = batch_as_tensor[0].shape[1]
    assert learner.normalize_batch(batch_as_tensor) == [
        (
            q + [NULL_TOKEN] * (q_len-len(q)),
            a
        )
        for (q, a) in batch
    ]
    
def test_Learner_learn(learner):
    with pytest.raises(RuntimeError, match="opt is None") as exc_info:
        learner.set_optimizer(None)
        learner.learn()
    learner.set_optimizer("AdamW")
    
    with pytest.raises(RuntimeError, match="expired") as exc_info:
        learner.add_words(["strange"], rebuild_optimizer=False)
        learner.learn()
    learner.remove_words(["strange"], rebuild_optimizer=True)

    batch = learner.make_batchs(4)[0]
    learner.learn(batch)

    # 勾配蓄積モードなら勾配が変化しない (パラメータは変化する)
    A_grad_old = learner.model.A[0].grad.detach().clone()
    A_old      = learner.model.A[0]     .detach().clone()
    learner.learn()
    A_grad_new = learner.model.A[0].grad
    A_new      = learner.model.A[0]
    assert_close(A_grad_old, A_grad_new)
    assert_not_close(A_old, A_new)

    # 通常モードなら勾配が変化する (当然パラメータも変化する)
    A_grad_old = learner.model.A[0].grad.detach().clone()
    A_old      = learner.model.A[0]     .detach().clone()
    learner.learn(batch)
    A_grad_new = learner.model.A[0].grad
    A_new      = learner.model.A[0]
    assert_not_close(A_grad_old, A_grad_new)
    assert_not_close(A_old, A_new)

    # 戻り値
    assert learner.learn() is None
    assert set(learner.learn          (batch).keys()) == \
           set(learner.accumulate_grad(batch).keys())

    # 通常モードで、バッチに未知語があれば、未知語が処理される
    learner.words_added = False
    def hook(*args, **kwargs):
        learner.words_added = True
    learner.on_word_added = hook
    learner.learn(strange="add")
    assert not learner.words_added
    learner.learn([(["未知語0"], "未知語1")], strange="add")
    assert learner.words_added

def test_Learner_zero_grad(learner):
    pass

def test_Learner_accumulate_grad(learner):
    batch = learner.make_batchs(5)[0]
    learner.learn(batch)

    # 勾配が変化する (パラメータは変化しない)
    A_grad_old = learner.model.A[0].grad.detach().clone()
    A_old      = learner.model.A[0]     .detach().clone()
    learner.accumulate_grad(batch)
    A_grad_new = learner.model.A[0].grad
    A_new      = learner.model.A[0]
    assert_not_close(A_grad_old, A_grad_new)
    assert_close(A_old, A_new)

    # 戻り値
    assert set(learner.accumulate_grad(batch)) == {
        "loss", "ce_loss", "norm_reg", "batch_size"
    }

    # 未知語が処理される
    learner.words_ignored = False
    def hook(*args, **kwargs):
        learner.words_ignored = True
    learner.on_strange_ignored = hook
    learner.accumulate_grad([(["未知語0"], "未知語1")], strange="ignore")
    assert learner.words_ignored

(テストコード test_project.py を表示 ここまで)

ComplexTaylorWordFunctionLM.py をテストしたときのテストコードに加筆して、

from ComplexTaylorWordFunctionLearner import ComplexTaylorWordFunctionLearner
と、 # Learner 以降を追加した。

上記のテストコードを用意した上で、カレントディレクトリにて pytest -s -v を実行した結果、次のようになった。

結果を表示
pytest -s -v
(cd)>pytest -s -v
================================================= test session starts =================================================
platform win32 -- Python 3.13.1, pytest-9.0.2, pluggy-1.6.0 -- C:\xxx\python.exe
cachedir: .pytest_cache
rootdir: (cd)
plugins: anyio-4.9.0
collected 29 items

test_project.py::test_Model_init_error PASSED
test_project.py::test_Model_init
model.A[0]=Parameter containing:
tensor([[[ 0.0000+0.0000j],
         [ 0.0000+0.0000j],
         [ 0.0000+0.0000j],
         [ 0.0000+0.0000j]],

        [[ 0.0048-0.0190j],
         [ 0.0254-0.0157j],
         [-0.0060+0.0150j],
         [-0.0185+0.0176j]],

        [[ 0.0125-0.0190j],
         [ 0.0051+0.0104j],
         [ 0.0115-0.0182j],
         [-0.0412+0.0213j]],

        [[-0.0243+0.0071j],
         [ 0.0206+0.0125j],
         [ 0.0408-0.0076j],
         [ 0.0509-0.0095j]],

        [[ 0.0115+0.0037j],
         [-0.0020+0.0157j],
         [-0.0336-0.0043j],
         [-0.0091+0.0266j]],

        [[ 0.0250+0.0259j],
         [ 0.0166+0.0082j],
         [ 0.0183-0.0047j],
         [-0.0180+0.0299j]],

        [[ 0.0051-0.0305j],
         [-0.0100-0.0014j],
         [-0.0236-0.0144j],
         [ 0.0192+0.0080j]]], requires_grad=True)
model.A[1]=Parameter containing:
tensor([[[ 1.0000e+00+0.0000j,  0.0000e+00+0.0000j,  0.0000e+00+0.0000j,
           0.0000e+00+0.0000j],
         [ 0.0000e+00+0.0000j,  1.0000e+00+0.0000j,  0.0000e+00+0.0000j,
           0.0000e+00+0.0000j],
         [ 0.0000e+00+0.0000j,  0.0000e+00+0.0000j,  1.0000e+00+0.0000j,
           0.0000e+00+0.0000j],
         [ 0.0000e+00+0.0000j,  0.0000e+00+0.0000j,  0.0000e+00+0.0000j,
           1.0000e+00+0.0000j]],

        [[-1.0955e-02-0.0056j, -1.6600e-02-0.0098j,  1.3548e-02+0.0068j,
           8.1916e-05+0.0088j],
         [-1.1258e-03-0.0018j,  8.0738e-03+0.0021j, -1.1794e-02-0.0049j,
          -3.3284e-03+0.0008j],
         [-4.1518e-03-0.0125j, -1.6607e-02+0.0018j,  1.2475e-02+0.0006j,
          -9.6314e-03-0.0031j],
         [-1.0194e-02+0.0063j, -4.9095e-03+0.0017j,  2.8294e-03+0.0034j,
           7.2892e-03-0.0084j]],

        [[ 5.8126e-03-0.0107j,  1.3837e-02+0.0017j,  1.6353e-03+0.0084j,
           1.2794e-02-0.0052j],
         [ 1.0617e-02-0.0054j, -1.4136e-02+0.0247j, -6.5637e-03+0.0062j,
           1.6619e-02+0.0072j],
         [-9.1959e-03-0.0145j, -5.6918e-03-0.0010j,  5.1161e-03-0.0112j,
          -2.8556e-03+0.0006j],
         [ 8.7933e-03-0.0016j,  5.4906e-03+0.0002j,  6.1732e-03-0.0209j,
           6.7767e-03-0.0106j]],

        [[ 2.9751e-03-0.0025j, -3.6132e-03-0.0077j,  1.6654e-03+0.0036j,
          -8.3002e-03+0.0045j],
         [-2.2255e-03+0.0061j, -3.4040e-03+0.0085j,  6.8251e-03+0.0189j,
           1.4850e-03+0.0173j],
         [ 1.6204e-03+0.0179j, -9.8987e-03+0.0144j, -7.0138e-05+0.0249j,
          -5.7304e-03+0.0105j],
         [-1.0917e-02-0.0125j, -1.6250e-02+0.0050j, -2.1917e-02-0.0017j,
           1.0651e-02+0.0105j]],

        [[ 3.6555e-03+0.0015j, -2.4830e-03-0.0085j,  6.0171e-04+0.0092j,
           3.4052e-03+0.0049j],
         [ 5.7991e-03+0.0058j, -3.3864e-04-0.0184j,  1.3141e-02-0.0067j,
          -1.1756e-02+0.0050j],
         [ 1.7392e-03-0.0169j, -4.0781e-03-0.0031j, -4.8354e-03+0.0039j,
           2.7869e-03-0.0257j],
         [ 9.0058e-03-0.0038j,  2.6461e-03+0.0048j, -1.2134e-02+0.0007j,
           1.9882e-03+0.0053j]],

        [[-5.8150e-03+0.0039j, -1.5769e-02-0.0043j, -7.3172e-03+0.0068j,
           5.8741e-03-0.0126j],
         [ 3.8978e-03+0.0015j, -3.2744e-03+0.0190j, -5.9625e-04+0.0074j,
          -1.3365e-02-0.0055j],
         [ 8.2660e-03-0.0066j,  9.7537e-03+0.0082j, -9.4992e-03-0.0010j,
          -4.7207e-03-0.0205j],
         [ 1.4221e-03-0.0012j,  1.6763e-03-0.0049j,  5.2098e-04-0.0073j,
           1.2441e-04+0.0220j]],

        [[-2.5821e-03-0.0080j, -1.0807e-02-0.0076j,  3.0686e-04-0.0061j,
           5.1232e-03+0.0048j],
         [-1.1061e-02+0.0010j,  3.1278e-03+0.0049j, -3.2097e-02-0.0155j,
           8.8525e-03-0.0007j],
         [ 3.7356e-04-0.0100j, -4.0169e-03-0.0169j,  5.6795e-03+0.0020j,
           8.5298e-04+0.0027j],
         [ 1.9934e-02-0.0109j,  2.4313e-03+0.0117j, -3.0281e-03-0.0109j,
           2.2230e-02-0.0043j]]], requires_grad=True)

model.Next[0]=Parameter containing:
tensor([[0.+0.j],
        [0.+0.j],
        [0.+0.j],
        [0.+0.j]], requires_grad=True)
model.Next[1]=Parameter containing:
tensor([[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]], requires_grad=True)

model.z0=tensor([0.3536+0.3536j, 0.3536+0.3536j, 0.3536+0.3536j, 0.3536+0.3536j])
PASSED
test_project.py::test_Model_reset_null PASSED
test_project.py::test_Model_tensor_power PASSED
test_project.py::test_Model_apply_word PASSED
test_project.py::test_Model_sequence_meaning PASSED
test_project.py::test_Model_apply_next PASSED
test_project.py::test_Model_word_meaning PASSED
test_project.py::test_Model_update_povm_cache PASSED
test_project.py::test_Model_probs_3_methods PASSED
test_project.py::test_Model_forward PASSED
test_project.py::test_Model_forward_sequence PASSED
test_project.py::test_Model_povm_normalization_error
err=9.406961112290446e-07
PASSED
test_project.py::test_Model_speak
(i=0,)model.speak(i, print_text=False)=''

(i=1,)model.speak(i, print_text=False)='<Null>'

(i=2,)model.speak(i, print_text=False)='you love'

(i=3,)model.speak(i, print_text=False)='much much <Null>'

(i=4,)model.speak(i, print_text=False)='much much <Null> <Null>'

(i=5,)model.speak(i, print_text=False)='much much very very you'

(i=6,)model.speak(i, print_text=False)='much thank love I thank very'

(i=7,)model.speak(i, print_text=False)='very I very <Null> very thank I'

(i=8,)model.speak(i, print_text=False)='thank thank much you love love love love'

(i=9,)model.speak(i, print_text=False)='thank thank thank <Null> <Null> much thank love very'

(i=0,)model.speak(i, return_list=True, print_text=False)=[]

(i=1,)model.speak(i, return_list=True, print_text=False)=['<Null>']

(i=2,)model.speak(i, return_list=True, print_text=False)=['thank', 'you']

(i=3,)model.speak(i, return_list=True, print_text=False)=['<Null>', 'much', 'love']

(i=4,)model.speak(i, return_list=True, print_text=False)=['thank', 'thank', '<Null>', 'thank']

(i=5,)model.speak(i, return_list=True, print_text=False)=['<Null>', 'you', 'you', '<Null>', 'you']

(i=6,)model.speak(i, return_list=True, print_text=False)=['much', 'much', 'very', 'very', 'much', 'love']

(i=7,)model.speak(i, return_list=True, print_text=False)=['<Null>', 'thank', 'you', 'you', 'you', '<Null>', 'I']

(i=8,)model.speak(i, return_list=True, print_text=False)=['thank', 'thank', 'much', 'much', 'very', 'you', 'very', 'thank']

(i=9,)model.speak(i, return_list=True, print_text=False)=['<Null>', 'thank', 'I', 'very', 'very', 'you', 'much', '<Null>', '<Null>']
PASSED
test_project.py::test_Model_operator_2_methods PASSED
test_project.py::test_Learner_init_error PASSED
test_project.py::test_Learner_init PASSED
test_project.py::test_Learner_set_optimizer PASSED
test_project.py::test_Learner_add_words PASSED
test_project.py::test_Learner_remove_words PASSED
test_project.py::test_Learner_make_examples PASSED
test_project.py::test_Learner_make_batchs PASSED
test_project.py::test_Learner_make_batch_tensor PASSED
test_project.py::test_Learner_update_cache PASSED
test_project.py::test_Learner_operator_l2_regularization PASSED
test_project.py::test_Learner_normalize_batch PASSED
test_project.py::test_Learner_learn PASSED
test_project.py::test_Learner_zero_grad PASSED
test_project.py::test_Learner_accumulate_grad PASSED

================================================= 29 passed in 3.46s ==================================================

(cd)>

(結果を表示 ここまで)

全体のステータスが (テスト数) passed in (経過時間)s

となっており、x failed がないので、テストはすべて成功である。

改めて、現在のステータスを表示しておく。

ディレクトリ構成とステータス
(カレントディレクトリ)
├ main.py                               # 😴メインプログラム
├ ComplexTaylorWordFunctionLM.py        # ✅モデルクラス
├ ComplexTaylorWordFunctionLearner.py   # ✅モデルを学習するためのクラス
├ constants.py                          # ✅複数ファイルで共有する定数など (コーパス以外)
├ corpus.py                             # 😴トイコーパス
├ ComplexTaylorWordFunctionEvaluator.py # 😴学習結果や経過を評価するためのクラス
└ test_project.py                       # ✅テスト用コード

ステータスの凡例
😴: 未着手
🛠️: 今から作成・編集する
📌: 新たに編集が必要になる
👀: 作成済み(未テスト)
✅: 作成済み (テスト成功済み or テストを実施しない)
⚠️: 問題あり

5. 次回予告

次回は、4-9節以降を書いていく。

モデルの評価をする ComplexToaylorWordFunctionEvaluator.py を書いて、実験の準備を進める。

↓次回の記事

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?