12fromlabmlimportexperiment13fromlabml.configsimportoption14fromlabml\_nn.transformersimportTransformerConfigs15fromlabml\_nn.transformers.configsimportFeedForwardConfigs16fromlabml\_nn.transformers.mlm.experimentimportTransformerMLM,ConfigsasMLMConfigs

19classConfigs(MLMConfigs):

29mix\_mlp:FeedForwardConfigs

32@option(Configs.mix\_mlp)33def\_mix\_mlp\_configs(c:Configs):

38conf=FeedForwardConfigs()

40conf.d\_model=c.seq\_len

42conf.activation='GELU'

45returnconf

48@option(Configs.transformer)49def\_transformer\_configs(c:Configs):

56conf=TransformerConfigs()

58conf.n\_src\_vocab=c.n\_tokens59conf.n\_tgt\_vocab=c.n\_tokens

61conf.d\_model=c.d\_model

63fromlabml\_nn.transformers.mlp\_mixerimportMLPMixer64conf.encoder\_attn=MLPMixer(c.mix\_mlp.ffn)

67returnconf

70defmain():

72experiment.create(name="mlp\_mixer\_mlm")

74conf=Configs()

76experiment.configs(conf,{

78'batch\_size':64,

81'seq\_len':32,

84'epochs':1024,

87'inner\_iterations':1,

90'd\_model':128,91'transformer.ffn.d\_ff':256,92'transformer.n\_heads':8,93'transformer.n\_layers':6,94'transformer.ffn.activation':'GELU',

97'mix\_mlp.d\_ff':128,

100'optimizer.optimizer':'Noam',101'optimizer.learning\_rate':1.,102})

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

108withexperiment.start():

110conf.run()

114if\_\_name\_\_=='\_\_main\_\_':115main()