16importtorch17fromtorchimportnn18fromtorch.utils.dataimportDataLoader,RandomSampler1920fromlabmlimportmonit,lab,tracker,experiment,logger21fromlabml.loggerimportText22fromlabml\_helpers.datasets.textimportTextFileDataset23fromlabml\_nn.optimizers.noamimportNoam24fromlabml\_nn.transformers.retroimportmodelasretro25fromlabml\_nn.transformers.retro.datasetimportDataset,RetroIndex26fromlabml\_nn.transformers.retro.modelimportRetroModel,NearestNeighborEncoder

29classSampler:

36def\_\_init\_\_(self,device:torch.device,model:retro.RetroModel,tds:TextFileDataset,chunk\_len:int):

43self.chunk\_len=chunk\_len44self.tds=tds45self.model=model46self.device=device

49self.index=RetroIndex()

51defretrieve\_nearest\_neighbours(self,chunk:str):

57neighbor\_offsets=self.index([chunk],None)

60text=self.tds.train61neighbors=[text[j:j+self.chunk\_len\*2]forjinneighbor\_offsets[0]]

64returnneighbors

66defsample(self,prompt:str,sample\_len:int):

72neighbors\_str=[]

75sampled=''

78foriinrange(sample\_len):

81whilelen(neighbors\_str)\<len(prompt)//self.chunk\_len:

83off=len(neighbors\_str)\*self.chunk\_len84chunk=prompt[off:off+self.chunk\_len]

86neighbors\_str.append(self.retrieve\_nearest\_neighbours(chunk))

89src=self.tds.text\_to\_i(prompt)

91neighbors=torch.stack([torch.stack([self.tds.text\_to\_i(n)forninchunk])forchunkinneighbors\_str])

94src=src.to(self.device)95neighbors=neighbors.to(self.device)

98res=self.model(src[None,:],neighbors[None,:,:,:])

101token=res[0,-1,:].argmax(dim=-1)

104prompt+=self.tds.itos[token.item()]105sampled+=self.tds.itos[token.item()]

108returnsampled

111classTrainer:

116def\_\_init\_\_(self,device:torch.device,model:retro.RetroModel,117dataloader:DataLoader,optimizer:torch.optim.Optimizer):

124self.optimizer=optimizer125self.device=device126self.dataloader=dataloader127self.model=model128self.loss\_func=nn.CrossEntropyLoss()

130def\_\_call\_\_(self):

136fori,(src,tgt,neighbors)inmonit.enum('Train',self.dataloader):

138src,tgt,neighbors=src.to(self.device),tgt.to(self.device),neighbors.to(self.device)

141res=self.model(src,neighbors)

143loss=self.loss\_func(res.view(-1,res.shape[-1]),tgt.view(-1))

146self.optimizer.zero\_grad()

148loss.backward()

150self.optimizer.step()

153tracker.save({'loss.train':loss})154tracker.add\_global\_step(len(src))

157deftrain():

163experiment.create(name='retro\_small')

166device=torch.device('cuda:0')

169tds=TextFileDataset(170lab.get\_data\_path()/'tiny\_shakespeare.txt',171list,172url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')

175train\_dataset=Dataset(lab.get\_data\_path()/'retro\_train\_dataset.json',tds)

178train\_dl=DataLoader(train\_dataset,179batch\_size=4,180sampler=RandomSampler(train\_dataset,replacement=True))

183chunk\_len=16184d\_model=128185d\_ff=512186n\_heads=16187d\_k=16

190nearest\_neighbor\_encoder=NearestNeighborEncoder(chunk\_len,6,{3},d\_model,n\_heads,d\_k,d\_ff)

192model=RetroModel(tds.n\_tokens,d\_model,6,193{3,5},194chunk\_len,n\_heads,d\_k,d\_ff,195encoder=nearest\_neighbor\_encoder)

197model=model.to(device)

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

201trainer=Trainer(device,model,train\_dl,optimizer)

203sampler=Sampler(device,model,tds,chunk\_len)

205prompt='''Second Citizen:\nOne word, good citizens.\n\nFirst Citizen:'''

208experiment.add\_pytorch\_models(model=model)

211withexperiment.start():

213forepochinmonit.loop(32):

215trainer()

217tracker.new\_line()

219logger.log([(prompt.replace('\n','\\n\n'),Text.subtle),220(sampler.sample(prompt,128).replace('\n','\\n\n'),Text.none)])

222experiment.save\_checkpoint()

226if\_\_name\_\_=='\_\_main\_\_':227train()