11fromtypingimportList1213importtorch14importtorch.nnasnn15fromlabml.loggerimportText1617fromlabmlimportexperiment,tracker,monit,logger18fromlabml.configsimportoption19fromlabml\_helpers.metrics.simple\_stateimportSimpleStateModule20fromlabml\_helpers.moduleimportModule21fromlabml\_helpers.train\_validimportBatchIndex,hook\_model\_outputs22fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs23fromlabml\_nn.transformers.xlimportTransformerXL,TransformerXLLayer

26classAutoregressiveModel(Module):

31def\_\_init\_\_(self,n\_vocab:int,d\_model:int,transformer:TransformerXL):32super().\_\_init\_\_()

34self.src\_embed=nn.Embedding(n\_vocab,d\_model)

36self.transformer=transformer

38self.generator=nn.Linear(d\_model,n\_vocab)

40self.mask\_x=None41self.mask\_mem=None

43defforward(self,x:torch.Tensor,mem:List[torch.Tensor]):

45m\_len=len(mem[0])ifmemelse0

47ifself.mask\_xisNoneorself.mask\_x.shape[0]\<len(x):48fromlabml\_nn.transformers.utilsimportsubsequent\_mask49self.mask\_x=subsequent\_mask(len(x)).to(x.device)

51ifself.mask\_memisNoneorself.mask\_mem.shape[1]\<m\_lenorself.mask\_mem.shape[0]\<len(x):52self.mask\_mem=self.mask\_x.new\_ones(len(x),m\_len,1)

55ifm\_len:56mask=torch.cat((self.mask\_mem[:len(x),:m\_len],self.mask\_x[:len(x),:len(x)]),dim=1)

58else:59mask=self.mask\_x[:len(x),:len(x)]

62x=self.src\_embed(x)

64res,mem=self.transformer(x,mem,mask)

66res=self.generator(res)

68returnres,mem

71classConfigs(NLPAutoRegressionConfigs):

78model:AutoregressiveModel

81d\_model:int=128

83heads:int=4

85dropout:float=0.0

87d\_ff:int=256

89n\_layers:int=6

91mem\_len:int=128

93memory=SimpleStateModule()

95definit(self):

97tracker.set\_scalar("accuracy.\*",True)98tracker.set\_scalar("loss.\*",True)

100hook\_model\_outputs(self.mode,self.model,'model')

102self.state\_modules=[self.accuracy,self.memory]

104defmerge\_memory(self,old\_mem,new\_mem):

111ifself.mem\_len==0:112return[]

115ifold\_mem:116mem=[torch.cat((m,x),dim=0)form,xinzip(old\_mem,new\_mem)]117else:118mem=new\_mem

121iflen(mem[0])\>self.mem\_len:122mem=[m[-self.mem\_len:]forminmem]

125returnmem

127defstep(self,batch:any,batch\_idx:BatchIndex):

133data,target=batch[0].to(self.device),batch[1].to(self.device)

136ifself.mode.is\_train:137tracker.add\_global\_step(data.shape[0]\*data.shape[1])

140withself.mode.update(is\_log\_activations=batch\_idx.is\_last):

142mem=self.memory.get()

144output,new\_mem=self.model(data,mem)

146mem=self.merge\_memory(mem,new\_mem)

148self.memory.set(mem)

151loss=self.loss\_func(output,target)152tracker.add("loss.",loss)

155self.accuracy(output,target)156self.accuracy.track()

159ifself.mode.is\_train:

161loss.backward()

163torch.nn.utils.clip\_grad\_norm\_(self.model.parameters(),max\_norm=self.grad\_norm\_clip)

165self.optimizer.step()

167ifbatch\_idx.is\_last:168tracker.add('model',self.model)

170self.optimizer.zero\_grad()

173tracker.save()

175defsample(self):

181prompt=self.prompt

183log=[(prompt,Text.subtle)]

185mem=[]

187foriinmonit.iterate('Sample',25):

189data=self.text.text\_to\_i(prompt).unsqueeze(-1)

191data=data.to(self.device)

193output,new\_mem=self.model(data,mem)

195output=output.argmax(dim=-1).squeeze(1)

197prompt+=self.prompt\_separator+self.text.itos[output[-1]]

199prompt=prompt[-1:]

201log+=[(self.prompt\_separator+self.text.itos[output[-1]],Text.value)]

203mem=self.merge\_memory(mem,new\_mem)

206logger.log(log)

209@option(Configs.model)210defautoregressive\_model(c:Configs):

214fromlabml\_nn.transformers.xlimportRelativeMultiHeadAttention215fromlabml\_nn.transformers.feed\_forwardimportFeedForward216m=AutoregressiveModel(c.n\_tokens,c.d\_model,TransformerXL(217TransformerXLLayer(d\_model=c.d\_model,218self\_attn=RelativeMultiHeadAttention(c.heads,c.d\_model,c.dropout),219feed\_forward=FeedForward(c.d\_model,c.d\_ff,c.dropout),220dropout\_prob=c.dropout),c.n\_layers))221returnm.to(c.device)

224defmain():

229experiment.create(name="transformer\_xl",comment='')

231conf=Configs()

233experiment.configs(conf,

235{'tokenizer':'character',236'text':'tiny\_shakespeare',237'optimizer.learning\_rate':1.,238'optimizer.optimizer':'Noam',239'prompt':'It is',240'prompt\_separator':'',241242'train\_loader':'sequential\_train\_loader',243'valid\_loader':'sequential\_valid\_loader',244245'seq\_len':2,246'mem\_len':32,247'epochs':128,248'batch\_size':32,249'inner\_iterations':25,250})

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

256withexperiment.start():

258conf.run()

262if\_\_name\_\_=='\_\_main\_\_':263main()