12importtorch13fromlabmlimportexperiment14fromlabml.configsimportoption15fromlabml.utils.pytorchimportget\_modules16fromlabml\_helpers.moduleimportModule1718fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs19fromlabml\_nn.transformersimportEncoder,Generator,TransformerConfigs20fromlabml\_nn.transformers.utilsimportsubsequent\_mask

23classAutoregressiveModel(Module):

28def\_\_init\_\_(self,src\_embed:Module,encoder:Encoder,generator:Generator,\*,29is\_save\_ff\_input:bool=False):30super().\_\_init\_\_()

32self.src\_embed=src\_embed

34self.encoder=encoder

38self.encoder.layers[-1].is\_save\_ff\_input=is\_save\_ff\_input

41self.generator=generator

43self.src\_mask=None

45@property46defff\_input(self)-\>torch.Tensor:

50returnself.encoder.layers[-1].ff\_input

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

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

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

59returnself.generator(res),None

62classConfigs(NLPAutoRegressionConfigs):

69transformer:TransformerConfigs70model:AutoregressiveModel7172is\_save\_ff\_input=False

75@option(Configs.model)76defautoregressive\_model(c:Configs):

80m=AutoregressiveModel(

83src\_embed=c.transformer.src\_embed,84encoder=c.transformer.encoder,85generator=c.transformer.generator,

87is\_save\_ff\_input=c.is\_save\_ff\_input)88returnm.to(c.device)

91@option(Configs.transformer)92deftransformer\_c(c:Configs):

96tc=TransformerConfigs()97tc.n\_src\_vocab=c.n\_tokens98tc.n\_tgt\_vocab=c.n\_tokens99100returntc

103defmain():

105experiment.create(name="knn\_lm")

107conf=Configs()

109experiment.configs(conf,

111{'tokenizer':'character',112'prompt\_separator':'',113'prompt':'It is ',114'text':'tiny\_shakespeare',115116'optimizer.optimizer':'Noam',117'optimizer.learning\_rate':1.,118'optimizer.d\_model':256,119120'seq\_len':1024,121'epochs':128,122'batch\_size':6,123'inner\_iterations':10,

126'transformer.d\_model':256,127'transformer.ffn.d\_ff':1024,128'transformer.n\_heads':8,129'transformer.n\_layers':6})

132conf.n\_tokens=conf.text.n\_tokens

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

138withexperiment.start():

140conf.run()141142143if\_\_name\_\_=='\_\_main\_\_':144main()