20importdataclasses2122importtorch23fromlabml\_helpers.moduleimportModule24fromtorchimportnn25fromtorch.utils.dataimportDataset,DataLoader2627fromlabmlimportexperiment,lab,tracker,monit,logger28fromlabml.loggerimportText29fromlabml.utils.downloadimportdownload\_file30fromlabml\_nn.experiments.nlp\_autoregressionimporttranspose\_batch31fromlabml\_nn.optimizers.noamimportNoam32fromlabml\_nn.transformersimportEncoder,MultiHeadAttention33fromlabml\_nn.transformers.feed\_forwardimportFeedForward34fromlabml\_nn.transformers.modelsimportEmbeddingsWithPositionalEncoding,TransformerLayer35fromlabml\_nn.transformers.utilsimportsubsequent\_mask

38classAutoregressiveModel(Module):

43def\_\_init\_\_(self,src\_embed:Module,encoder:Encoder,generator:Module):44super().\_\_init\_\_()

46self.src\_embed=src\_embed

48self.encoder=encoder

51self.generator=generator

53self.src\_mask=None

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

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

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

62returnself.generator(res)

[email protected]:

70d\_model:int=51271seq\_len:int=12872batch\_size:int=3273n\_layers:int=674n\_heads:int=875dropout:float=0.176d\_ff:int=204877glu\_variant:str='GLU'78epochs:int=579grad\_norm\_clip:float=0.5

82classTinyShakespeareDataset(Dataset):

87def\_\_init\_\_(self,seq\_len:int):

89path=lab.get\_data\_path()/'tiny\_shakespeare.txt'

91download\_file('https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt',path)

93withopen(str(path),'r')asf:94text=f.read()

97chars=list(set(text))

99self.stoi={c:ifori,cinenumerate(chars)}

101self.itos={i:cfori,cinenumerate(chars)}

103self.seq\_len=seq\_len

105self.data=self.text\_to\_i(text)

107deftext\_to\_i(self,text:str):

111returntorch.tensor([self.stoi[c]forcintext],dtype=torch.long)

113def\_\_len\_\_(self):

119returnlen(self.data)-self.seq\_len-1

121def\_\_getitem\_\_(self,idx):

125returnself.data[idx:idx+self.seq\_len],self.data[idx+1:idx+self.seq\_len+1]

128classTrainer:

133def\_\_init\_\_(self,configs:Configs):

135self.device=torch.device('cpu')136iftorch.cuda.is\_available():137self.device=torch.device('cuda:0')

139self.dataset=TinyShakespeareDataset(configs.seq\_len)

141self.dataloader=DataLoader(self.dataset,142batch\_size=configs.batch\_size,143collate\_fn=transpose\_batch,144shuffle=True)

148ifconfigs.glu\_variant=='GLU':149ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.Sigmoid(),True,False,False,False)

152elifconfigs.glu\_variant=='Bilinear':153ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.Identity(),True,False,False,False)

156elifconfigs.glu\_variant=='ReGLU':157ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.ReLU(),True,False,False,False)

160elifconfigs.glu\_variant=='GEGLU':161ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.GELU(),True,False,False,False)

165elifconfigs.glu\_variant=='SwiGLU':166ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.SiLU(),True,False,False,False)

169elifconfigs.glu\_variant=='ReLU':170ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.ReLU())

173elifconfigs.glu\_variant=='GELU':174ffn=FeedForward(configs.d\_model,configs.d\_ff,configs.dropout,nn.GELU())175else:176raiseValueError(f'Unknown variant {configs.glu\_variant}')

179n\_chars=len(self.dataset.stoi)

182mha=MultiHeadAttention(configs.n\_heads,configs.d\_model,configs.dropout)

184transformer\_layer=TransformerLayer(d\_model=configs.d\_model,self\_attn=mha,src\_attn=None,185feed\_forward=ffn,dropout\_prob=configs.dropout)

191self.model=AutoregressiveModel(EmbeddingsWithPositionalEncoding(configs.d\_model,n\_chars),192Encoder(transformer\_layer,configs.n\_layers),193nn.Linear(configs.d\_model,n\_chars))

196self.model.to(self.device)

199self.optimizer=Noam(self.model.parameters(),lr=1.0,warmup=2\_000,d\_model=configs.d\_model)

202self.loss\_func=nn.CrossEntropyLoss()

205self.epochs=configs.epochs

207self.grad\_norm\_clip=configs.grad\_norm\_clip

210tracker.set\_scalar("loss.\*",True)

212defsample(self):

218prompt='It is'

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

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

224data=self.dataset.text\_to\_i(prompt).unsqueeze(-1)225data=data.to(self.device)

227output=self.model(data)

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

231prompt+=self.dataset.itos[output[-1].item()]

233log+=[(self.dataset.itos[output[-1].item()],Text.value)]

236logger.log(log)

238deftrain(self):

244for\_inmonit.loop(self.epochs):

246fori,batchinmonit.enum('Train',self.dataloader):

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

251tracker.add\_global\_step(data.shape[0]\*data.shape[1])

254self.model.train()

256output=self.model(data)

259loss=self.loss\_func(output.view(-1,output.shape[-1]),target.view(-1))

261tracker.add("loss.train",loss)

264loss.backward()

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

268self.optimizer.step()

270if(i+1)%100==0:271tracker.add('model',self.model)

273self.optimizer.zero\_grad()

276if(i+1)%100==0:277self.model.eval()278withtorch.no\_grad():279self.sample()

282if(i+1)%10==0:283tracker.save()

286experiment.save\_checkpoint()

289defmain():

291experiment.create(name="glu\_variants")

293configs=Configs()

295experiment.configs(dataclasses.asdict(configs))

298trainer=Trainer(configs)

300experiment.add\_pytorch\_models({'model':trainer.model})

303withexperiment.start():

305trainer.train()306307308if\_\_name\_\_=='\_\_main\_\_':309main()