11fromtypingimportList1213importtorch14fromtorchimportnn1516fromlabmlimportexperiment,tracker,logger17fromlabml.configsimportoption18fromlabml.loggerimportText19fromlabml\_helpers.metrics.accuracyimportAccuracy20fromlabml\_helpers.moduleimportModule21fromlabml\_helpers.train\_validimportBatchIndex22fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs23fromlabml\_nn.transformersimportEncoder,Generator24fromlabml\_nn.transformersimportTransformerConfigs25fromlabml\_nn.transformers.mlmimportMLM

28classTransformerMLM(nn.Module):

33def\_\_init\_\_(self,\*,encoder:Encoder,src\_embed:Module,generator:Generator):

40super().\_\_init\_\_()41self.generator=generator42self.src\_embed=src\_embed43self.encoder=encoder

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

47x=self.src\_embed(x)

49x=self.encoder(x,None)

51y=self.generator(x)

55returny,None

58classConfigs(NLPAutoRegressionConfigs):

69model:TransformerMLM

71transformer:TransformerConfigs

74n\_tokens:int='n\_tokens\_mlm'

76no\_mask\_tokens:List[int]=[]

78masking\_prob:float=0.15

80randomize\_prob:float=0.1

82no\_change\_prob:float=0.1

84mlm:MLM

87mask\_token:int

89padding\_token:int

92prompt:str=[93"We are accounted poor citizens, the patricians good.",94"What authority surfeits on would relieve us: if they",95"would yield us but the superfluity, while it were",96"wholesome, we might guess they relieved us humanely;",97"but they think we are too dear: the leanness that",98"afflicts us, the object of our misery, is as an",99"inventory to particularise their abundance; our",100"sufferance is a gain to them Let us revenge this with",101"our pikes, ere we become rakes: for the gods know I",102"speak this in hunger for bread, not in thirst for revenge.",103]

105definit(self):

111self.mask\_token=self.n\_tokens-1

113self.padding\_token=self.n\_tokens-2

116self.mlm=MLM(padding\_token=self.padding\_token,117mask\_token=self.mask\_token,118no\_mask\_tokens=self.no\_mask\_tokens,119n\_tokens=self.n\_tokens,120masking\_prob=self.masking\_prob,121randomize\_prob=self.randomize\_prob,122no\_change\_prob=self.no\_change\_prob)

125self.accuracy=Accuracy(ignore\_index=self.padding\_token)

127self.loss\_func=nn.CrossEntropyLoss(ignore\_index=self.padding\_token)

129super().init()

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

137data=batch[0].to(self.device)

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

144withtorch.no\_grad():145data,labels=self.mlm(data)

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

152output,\*\_=self.model(data)

155loss=self.loss\_func(output.view(-1,output.shape[-1]),labels.view(-1))156tracker.add("loss.",loss)

159self.accuracy(output,labels)160self.accuracy.track()

163ifself.mode.is\_train:

165loss.backward()

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

169self.optimizer.step()

171ifbatch\_idx.is\_last:172tracker.add('model',self.model)

174self.optimizer.zero\_grad()

177tracker.save()

[email protected]\_grad()180defsample(self):

186data=torch.full((self.seq\_len,len(self.prompt)),self.padding\_token,dtype=torch.long)

188fori,pinenumerate(self.prompt):

190d=self.text.text\_to\_i(p)

192s=min(self.seq\_len,len(d))193data[:s,i]=d[:s]

195data=data.to(self.device)

198data,labels=self.mlm(data)

200output,\*\_=self.model(data)

203forjinrange(data.shape[1]):

205log=[]

207foriinrange(len(data)):

209iflabels[i,j]!=self.padding\_token:

211t=output[i,j].argmax().item()

213ift\<len(self.text.itos):

215ift==labels[i,j]:216log.append((self.text.itos[t],Text.value))

218else:219log.append((self.text.itos[t],Text.danger))

221else:222log.append(('\*',Text.danger))

224elifdata[i,j]\<len(self.text.itos):225log.append((self.text.itos[data[i,j]],Text.subtle))

228logger.log(log)

231@option(Configs.n\_tokens)232defn\_tokens\_mlm(c:Configs):

236returnc.text.n\_tokens+2

239@option(Configs.transformer)240def\_transformer\_configs(c:Configs):

247conf=TransformerConfigs()

249conf.n\_src\_vocab=c.n\_tokens250conf.n\_tgt\_vocab=c.n\_tokens

252conf.d\_model=c.d\_model

255returnconf

258@option(Configs.model)259def\_model(c:Configs):

263m=TransformerMLM(encoder=c.transformer.encoder,264src\_embed=c.transformer.src\_embed,265generator=c.transformer.generator).to(c.device)266267returnm

270defmain():

272experiment.create(name="mlm")

274conf=Configs()

276experiment.configs(conf,{

278'batch\_size':64,

281'seq\_len':32,

284'epochs':1024,

287'inner\_iterations':1,

290'd\_model':128,291'transformer.ffn.d\_ff':256,292'transformer.n\_heads':8,293'transformer.n\_layers':6,

296'optimizer.optimizer':'Noam',297'optimizer.learning\_rate':1.,298})

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

304withexperiment.start():

306conf.run()

310if\_\_name\_\_=='\_\_main\_\_':311main()