16importtorch17fromtorchimportnn1819fromlabmlimportexperiment20fromlabml.configsimportoption21fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs22fromlabml\_nn.transformersimportTransformerConfigs,Encoder23fromlabml\_nn.transformers.utilsimportsubsequent\_mask

26classAutoregressiveTransformer(nn.Module):

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

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

43self.mask=None

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

48ifself.maskisNoneorself.mask.size(0)!=len(x):

50self.mask=subsequent\_mask(len(x)).to(x.device)

52x=self.src\_embed(x)

54x=self.encoder(x,self.mask)

56x=self.generator(x)

60returnx,None

63classConfigs(NLPAutoRegressionConfigs):

72model:AutoregressiveTransformer

74transformer:TransformerConfigs

77@option(Configs.transformer,'Transformer')78def\_transformer\_configs(c:Configs):

85conf=TransformerConfigs()

87conf.n\_src\_vocab=c.n\_tokens88conf.n\_tgt\_vocab=c.n\_tokens

90conf.d\_model=c.d\_model

93returnconf

96@option(Configs.model)97def\_model(c:Configs):

101m=AutoregressiveTransformer(c.transformer.encoder,102c.transformer.src\_embed,103c.transformer.generator).to(c.device)104105returnm

108defmain():

110experiment.create(name="transformer")

112conf=Configs()

114experiment.configs(conf,{

116'tokenizer':'character',

118'prompt\_separator':'',

120'prompt':'It is ',

122'text':'tiny\_shakespeare',

125'seq\_len':512,

127'epochs':32,

129'batch\_size':16,

132'inner\_iterations':10,

135'd\_model':256,136'transformer.n\_heads':16,137'transformer.ffn.d\_ff':1024,

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