18fromlabmlimportexperiment19fromlabml.configsimportoption20fromlabml\_nn.transformersimportTransformerConfigs21fromlabml\_nn.transformers.basic.autoregressive\_experimentimportConfigsasBasicAutoRegressionConfigs22fromlabml\_nn.transformers.gmlpimportGMLPBlock

25classConfigs(BasicAutoRegressionConfigs):

34transformer:TransformerConfigs='gMLP'

36gmlp:GMLPBlock

38d\_ffn:int=2048

41@option(Configs.gmlp,'gMLP')42def\_gmlp\_configs(c:Configs):

46returnGMLPBlock(c.d\_model,c.d\_ffn,c.seq\_len)

49@option(Configs.transformer,'gMLP')50def\_transformer\_configs(c:Configs):

57conf=TransformerConfigs()

59conf.n\_src\_vocab=c.n\_tokens60conf.n\_tgt\_vocab=c.n\_tokens

62conf.d\_model=c.d\_model

64conf.encoder\_layer=c.gmlp6566returnconf

69defmain():

71experiment.create(name="gMLP")

73conf=Configs()

75experiment.configs(conf,{

77'tokenizer':'character',

79'prompt\_separator':'',

81'prompt':'It is ',

83'text':'tiny\_shakespeare',

86'seq\_len':256,

88'epochs':128,

90'batch\_size':32,

93'inner\_iterations':10,

96'd\_model':512,97'd\_ffn':2048,

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()