19importtorch20fromtorchimportnn2122fromlabmlimportexperiment23fromlabml.configsimportoption24fromlabml.utils.pytorchimportget\_modules25fromlabml\_helpers.moduleimportModule2627fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs28fromlabml\_nn.transformersimportEncoder,Generator,TransformerConfigs29fromlabml\_nn.transformers.utilsimportsubsequent\_mask

32classAutoregressiveModel(Module):

37def\_\_init\_\_(self,n\_vocab:int,d\_model:int,transformer:Module):38super().\_\_init\_\_()

40self.src\_embed=nn.Embedding(n\_vocab,d\_model)41self.transformer=transformer42self.generator=nn.Linear(d\_model,n\_vocab)

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

46x=self.src\_embed(x)

48res=self.transformer(x)

50returnself.generator(res),None

53classConfigs(NLPAutoRegressionConfigs):

60model:AutoregressiveModel6162d\_model:int=51263heads:int=864dropout:float=0.065d\_ff:int=204866n\_layers:int=6

69@option(Configs.model)70deffeedback\_transformer(c:Configs):

74fromlabml\_nn.transformers.feedbackimportFeedbackTransformer,FeedbackTransformerLayer,\75FeedbackAttention,FeedForward7677returnAutoregressiveModel(78c.n\_tokens,c.d\_model,79FeedbackTransformer(80FeedbackTransformerLayer(d\_model=c.d\_model,81attn=FeedbackAttention(c.heads,c.d\_model,c.dropout),82feed\_forward=FeedForward(c.d\_model,c.d\_ff,c.dropout),83dropout\_prob=c.dropout),84c.n\_layers)).to(c.device)

87@option(Configs.model)88deffeedback\_transformer\_kv(c:Configs):

92fromlabml\_nn.transformers.feedbackimportFeedbackTransformerKV,FeedbackTransformerLayer,\93FeedbackAttention,FeedForward9495returnAutoregressiveModel(96c.n\_tokens,c.d\_model,97FeedbackTransformerKV(98FeedbackTransformerLayer(d\_model=c.d\_model,99attn=FeedbackAttention(c.heads,c.d\_model,c.dropout,100is\_kv\_precomputed=True),101feed\_forward=FeedForward(c.d\_model,c.d\_ff,c.dropout),102dropout\_prob=c.dropout),103c.n\_layers,c.d\_model,c.heads)).to(c.device)

106defmain():

108experiment.create(name="feedback\_transformer")

110conf=Configs()

112experiment.configs(conf,

114{'tokenizer':'character',115'text':'tiny\_shakespeare',116'optimizer.learning\_rate':1.0,117'optimizer.optimizer':'Noam',118'prompt':'It is',119'prompt\_separator':'',

122'model':'feedback\_transformer\_kv',123124'train\_loader':'shuffled\_train\_loader',125'valid\_loader':'shuffled\_valid\_loader',126127'seq\_len':128,128'epochs':128,129'batch\_size':64,130'inner\_iterations':25})

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

136withexperiment.start():

138conf.run()139140141if\_\_name\_\_=='\_\_main\_\_':142main()