#

列車の自己回帰変圧器

12import torch
13from labml import experiment
14from labml.configs import option
15from labml.utils.pytorch import get_modules
16from labml_helpers.module import Module
17
18from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
19from labml_nn.transformers import Encoder, Generator, TransformerConfigs
20from labml_nn.transformers.utils import subsequent_mask

#

自動回帰モデル

23class AutoregressiveModel(Module):

#

28    def __init__(self, src_embed: Module, encoder: Encoder, generator: Generator, *,
29                 is_save_ff_input: bool = False):
30        super().__init__()

#

トークン埋め込みモジュール

32        self.src_embed = src_embed

#

トランスベースのエンコーダ

34        self.encoder = encoder

#

エンコーダーの最後の層で入力をフィードフォワード層に保存するかどうか。これで出ました $f (c_{t})$ 。コンテキストの埋め込みです

。

38        self.encoder.layers[-1].is_save_ff_input = is_save_ff_input

#

次のトークン生成レイヤー。これにより、次のトークンのロジットが返されます

41        self.generator = generator

#

これは最初の呼び出しで初期化されます。

43        self.src_mask = None

#

検索保存済み $f (c_{t})$

45    @property
46    def ff_input(self) -> torch.Tensor:

#

50        return self.encoder.layers[-1].ff_input

#

52    def forward(self, src: torch.Tensor):

#

次のマスクを作成して、トランスフォーマーが過去のトークンにしか注目できないようにします。

54        if self.src_mask is None or self.src_mask.size(0) != len(src):
55            self.src_mask = subsequent_mask(len(src)).to(src.device)

#

トークン (src ) を埋め込み、トランスフォーマーに通します

57        res = self.encoder(self.src_embed(src), self.src_mask)

#

次のトークンのロジットを生成

59        return self.generator(res), None

#

コンフィギュレーション

デフォルトの設定は、実験を開始したときに上書きでき、また上書きされます。

62class Configs(NLPAutoRegressionConfigs):

#

69    transformer: TransformerConfigs
70    model: AutoregressiveModel
71
72    is_save_ff_input = False

#

自己回帰モデルを初期化

75@option(Configs.model)
76def autoregressive_model(c: Configs):

#

80    m = AutoregressiveModel(

#

設定可能なトランスフォーマーからソーストークン埋め込みレイヤー、エンコーダー、最終トークンジェネレーターを取得

83        src_embed=c.transformer.src_embed,
84        encoder=c.transformer.encoder,
85        generator=c.transformer.generator,

#

保存するかどうか $f (c_{t})$

87        is_save_ff_input=c.is_save_ff_input)
88    return m.to(c.device)

#

自己回帰モデルの設定可能なトランスエンコーダーを初期化します。

91@option(Configs.transformer)
92def transformer_c(c: Configs):

#

96    tc = TransformerConfigs()
97    tc.n_src_vocab = c.n_tokens
98    tc.n_tgt_vocab = c.n_tokens
99
100    return tc

#

103def main():

#

実験を作成

105    experiment.create(name="knn_lm")

#

コンフィグの作成

107    conf = Configs()

#

構成をロード

109    experiment.configs(conf,

#

オーバーライドする設定の辞書

111                       {'tokenizer': 'character',
112                        'prompt_separator': '',
113                        'prompt': 'It is ',
114                        'text': 'tiny_shakespeare',
115
116                        'optimizer.optimizer': 'Noam',
117                        'optimizer.learning_rate': 1.,
118                        'optimizer.d_model': 256,
119
120                        'seq_len': 1024,
121                        'epochs': 128,
122                        'batch_size': 6,
123                        'inner_iterations': 10,

#

変圧器構成

126                        'transformer.d_model': 256,
127                        'transformer.ffn.d_ff': 1024,
128                        'transformer.n_heads': 8,
129                        'transformer.n_layers': 6})

#

これはモデルを初期化するために必要です

132    conf.n_tokens = conf.text.n_tokens

#

保存および読み込み用のモデルを設定する

135    experiment.add_pytorch_models(get_modules(conf))

#

実験を始める

138    with experiment.start():

#

TrainValidConfigs.run

140        conf.run()
141
142
143if __name__ == '__main__':
144    main()