16importtorch17fromlabmlimportexperiment18fromlabml.configsimportoption19fromlabml.utils.pytorchimportget\_modules20fromlabml\_helpers.moduleimportModule2122fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs23fromlabml\_nn.transformersimportEncoder,Generator,TransformerConfigs24fromlabml\_nn.transformers.utilsimportsubsequent\_mask

27classAutoregressiveModel(Module):

32def\_\_init\_\_(self,src\_embed:Module,encoder:Encoder,generator:Generator):33super().\_\_init\_\_()

35self.src\_embed=src\_embed

37self.encoder=encoder

40self.generator=generator

42self.src\_mask=None

44defforward(self,src:torch.Tensor):

46ifself.src\_maskisNoneorself.src\_mask.size(0)!=len(src):47self.src\_mask=subsequent\_mask(len(src)).to(src.device)

49res=self.encoder(self.src\_embed(src),self.src\_mask)

51returnself.generator(res),None

54classConfigs(NLPAutoRegressionConfigs):

61transformer:TransformerConfigs62model:AutoregressiveModel

65@option(Configs.model)66defautoregressive\_model(c:Configs):

70m=AutoregressiveModel(c.transformer.src\_embed,c.transformer.encoder,c.transformer.generator)71returnm.to(c.device)

74@option(Configs.transformer)75deftransformer\_c(c:Configs):

79tc=TransformerConfigs()80tc.n\_src\_vocab=c.n\_tokens81tc.n\_tgt\_vocab=c.n\_tokens8283returntc

86defmain():

88experiment.create(name="glu\_variants")

90conf=Configs()

92experiment.configs(conf,

94{'tokenizer':'character',95'prompt\_separator':'',96'prompt':'It is ',97'text':'tiny\_shakespeare',9899'optimizer.optimizer':'Noam',100'optimizer.learning\_rate':1.,101'optimizer.d\_model':256,102103'seq\_len':1024,104'epochs':128,105'batch\_size':6,106'inner\_iterations':10,

112'transformer.ffn.glu\_variant':'Bilinear',

115'transformer.d\_model':256,116'transformer.ffn.d\_ff':1024,117'transformer.n\_heads':8,118'transformer.n\_layers':6})

121conf.n\_tokens=conf.text.n\_tokens

124experiment.add\_pytorch\_models(get\_modules(conf))

127withexperiment.start():

129conf.run()130131132if\_\_name\_\_=='\_\_main\_\_':133main()