15import torch
16from torch import nn
17
18from labml import experiment
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_nn.experiments.nlp_classification import NLPClassificationConfigs
22from labml_nn.transformers import Encoder
23from labml_nn.transformers import TransformerConfigs

26class TransformerClassifier(nn.Module):

30    def __init__(self, encoder: Encoder, src_embed: Module, generator: nn.Linear):

37        super().__init__()
38        self.src_embed = src_embed
39        self.encoder = encoder
40        self.generator = generator

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

44        x = self.src_embed(x)

46        x = self.encoder(x, None)

52        x = self.generator(x[-1])

56        return x, None

59class Configs(NLPClassificationConfigs):

68    model: TransformerClassifier

70    transformer: TransformerConfigs

73@option(Configs.transformer)
74def _transformer_configs(c: Configs):

81    conf = TransformerConfigs()

83    conf.n_src_vocab = c.n_tokens
84    conf.n_tgt_vocab = c.n_tokens

87    return conf

90@option(TransformerConfigs.encoder_attn)
91def fnet_mix():

97    from labml_nn.transformers.fnet import FNetMix
98    return FNetMix()

101@option(Configs.model)
102def _model(c: Configs):

106    m = TransformerClassifier(c.transformer.encoder,
107                              c.transformer.src_embed,
108                              nn.Linear(c.d_model, c.n_classes)).to(c.device)
109
110    return m

113def main():

115    experiment.create(name="fnet")

117    conf = Configs()

119    experiment.configs(conf, {

121        'tokenizer': 'basic_english',

124        'epochs': 32,

127        'inner_iterations': 10,

130        'transformer.d_model': 512,
131        'transformer.ffn.d_ff': 2048,
132        'transformer.n_heads': 8,
133        'transformer.n_layers': 6,

137        'transformer.encoder_attn': 'fnet_mix',

140        'optimizer.optimizer': 'Noam',
141        'optimizer.learning_rate': 1.,
142    })

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

148    with experiment.start():

150        conf.run()

154if __name__ == '__main__':
155    main()