අවධානය නිදහස් ට්රාන්ස්ෆෝමර් (AFT) අත්හදා බැලීම

16import torch
17
18from labml import experiment
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
22from labml_nn.transformers import TransformerConfigs, Encoder
23from labml_nn.transformers.utils import subsequent_mask

26class AutoregressiveTransformer(Module):

34    def __init__(self, encoder: Encoder, src_embed: Module, generator: Module):

41        super().__init__()
42        self.src_embed = src_embed
43        self.encoder = encoder
44        self.generator = generator

47        self.mask = None

49    def forward(self, x: torch.Tensor):

52        if self.mask is None or self.mask.size(0) != len(x):

54            self.mask = subsequent_mask(len(x)).to(x.device)

57        x = self.src_embed(x)

59        x = self.encoder(x, self.mask)

61        x = self.generator(x)

65        return x, None

68class Configs(NLPAutoRegressionConfigs):

77    model: AutoregressiveTransformer

79    transformer: TransformerConfigs
80
81    local_window_size: int = 32

84@option(Configs.transformer, 'Transformer')
85def _transformer_configs(c: Configs):

92    conf = TransformerConfigs()

94    conf.n_src_vocab = c.n_tokens
95    conf.n_tgt_vocab = c.n_tokens

97    conf.d_model = c.d_model

99    from labml_nn.transformers.aft import AFTLocal
100    conf.encoder_attn = AFTLocal(c.d_model, c.seq_len, c.local_window_size)

103    return conf

106@option(Configs.model)
107def _model(c: Configs):

111    m = AutoregressiveTransformer(c.transformer.encoder,
112                                  c.transformer.src_embed,
113                                  c.transformer.generator).to(c.device)
114
115    return m

118def main():

120    experiment.create(name="aft")

122    conf = Configs()

124    experiment.configs(conf, {

126        'tokenizer': 'character',

128        'prompt_separator': '',

130        'prompt': 'It is ',

132        'text': 'tiny_shakespeare',

135        'seq_len': 256,

137        'epochs': 128,

139        'batch_size': 32,

142        'inner_iterations': 10,

145        'd_model': 128,

147        'transformer.ffn.d_ff': 256,

150        'optimizer.optimizer': 'Noam',
151        'optimizer.learning_rate': 1.,
152    })

155    experiment.add_pytorch_models({'model': conf.model})

158    with experiment.start():

160        conf.run()

164if __name__ == '__main__':
165    main()