16importtorch1718fromlabmlimportexperiment19fromlabml.configsimportoption20fromlabml\_helpers.moduleimportModule21fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs22fromlabml\_nn.transformersimportTransformerConfigs,Encoder23fromlabml\_nn.transformers.utilsimportsubsequent\_mask

26classAutoregressiveTransformer(Module):

34def\_\_init\_\_(self,encoder:Encoder,src\_embed:Module,generator:Module):

41super().\_\_init\_\_()42self.src\_embed=src\_embed43self.encoder=encoder44self.generator=generator

47self.mask=None

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

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

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

57x=self.src\_embed(x)

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

61x=self.generator(x)

65returnx,None

68classConfigs(NLPAutoRegressionConfigs):

77model:AutoregressiveTransformer

79transformer:TransformerConfigs8081local\_window\_size:int=32

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

92conf=TransformerConfigs()

94conf.n\_src\_vocab=c.n\_tokens95conf.n\_tgt\_vocab=c.n\_tokens

97conf.d\_model=c.d\_model

99fromlabml\_nn.transformers.aftimportAFTLocal100conf.encoder\_attn=AFTLocal(c.d\_model,c.seq\_len,c.local\_window\_size)

103returnconf

106@option(Configs.model)107def\_model(c:Configs):

111m=AutoregressiveTransformer(c.transformer.encoder,112c.transformer.src\_embed,113c.transformer.generator).to(c.device)114115returnm

118defmain():

120experiment.create(name="aft")

122conf=Configs()

124experiment.configs(conf,{

126'tokenizer':'character',

128'prompt\_separator':'',

130'prompt':'It is ',

132'text':'tiny\_shakespeare',

135'seq\_len':256,

137'epochs':128,

139'batch\_size':32,

142'inner\_iterations':10,

145'd\_model':128,

147'transformer.ffn.d\_ff':256,

150'optimizer.optimizer':'Noam',151'optimizer.learning\_rate':1.,152})

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

158withexperiment.start():

160conf.run()

164if\_\_name\_\_=='\_\_main\_\_':165main()