15importtorch16fromtorchimportnn1718fromlabmlimportexperiment19fromlabml.configsimportoption20fromlabml\_helpers.moduleimportModule21fromlabml\_nn.experiments.nlp\_classificationimportNLPClassificationConfigs22fromlabml\_nn.transformersimportEncoder23fromlabml\_nn.transformersimportTransformerConfigs

26classTransformerClassifier(nn.Module):

30def\_\_init\_\_(self,encoder:Encoder,src\_embed:Module,generator:nn.Linear):

37super().\_\_init\_\_()38self.src\_embed=src\_embed39self.encoder=encoder40self.generator=generator

42defforward(self,x:torch.Tensor):

44x=self.src\_embed(x)

46x=self.encoder(x,None)

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

56returnx,None

59classConfigs(NLPClassificationConfigs):

68model:TransformerClassifier

70transformer:TransformerConfigs

73@option(Configs.transformer)74def\_transformer\_configs(c:Configs):

81conf=TransformerConfigs()

83conf.n\_src\_vocab=c.n\_tokens84conf.n\_tgt\_vocab=c.n\_tokens

87returnconf

90@option(TransformerConfigs.encoder\_attn)91deffnet\_mix():

97fromlabml\_nn.transformers.fnetimportFNetMix98returnFNetMix()

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

106m=TransformerClassifier(c.transformer.encoder,107c.transformer.src\_embed,108nn.Linear(c.d\_model,c.n\_classes)).to(c.device)109110returnm

113defmain():

115experiment.create(name="fnet")

117conf=Configs()

119experiment.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})

145experiment.add\_pytorch\_models({'model':conf.model})

148withexperiment.start():

150conf.run()

154if\_\_name\_\_=='\_\_main\_\_':155main()