15importjson16frompathlibimportPath1718importnumpyasnp19importtorch20fromtorch.utils.dataimportDatasetasPyTorchDataset2122fromlabmlimportlab,monit23fromlabml\_helpers.datasets.textimportTextFileDataset,TextDataset24fromlabml\_nn.transformers.retro.databaseimportRetroIndex

27defbuild\_dataset(chunk\_len:int=16,chunks\_per\_sample:int=32,skip\_range:int=8):

39dataset=TextFileDataset(40lab.get\_data\_path()/'tiny\_shakespeare.txt',41list,42url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')

45text=dataset.train

48index=RetroIndex()

51sample\_offsets=[]

53i=054whilei\<len(text):

56skip=np.random.randint(skip\_range)57i+=skip

60ifi+chunks\_per\_sample\*chunk\_len\>len(text):61break

64sample\_offsets.append(i)

67i+=chunks\_per\_sample\*chunk\_len

70samples=[]

72foriinmonit.iterate('Gather Neighbors',sample\_offsets):

74sample=text[i:i+chunks\_per\_sample\*chunk\_len+1]

76src=sample[:-1]

78chunks=[src[j:j+chunk\_len]forjinrange(0,len(src),chunk\_len)]

80chunk\_offsets=[j+iforjinrange(0,len(src),chunk\_len)]

83neighbor\_offsets=index(chunks,chunk\_offsets)

86neighbors=[[text[j:j+chunk\_len\*2]forjinn\_off]forn\_offinneighbor\_offsets]

89samples.append((sample[:-1],sample[1:],neighbors))

94withopen(str(lab.get\_data\_path()/'retro\_train\_dataset.json'),'w')asf:95f.write(json.dumps(samples))

98classDataset(PyTorchDataset):

105def\_\_init\_\_(self,file\_path:Path,tds:TextDataset):

111self.tds=tds

113withopen(str(file\_path),'r')asf:114self.samples=json.loads(f.read())

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

120returnlen(self.samples)

122def\_\_getitem\_\_(self,idx:int):

127s=self.samples[idx]

129src=self.tds.text\_to\_i(s[0])130tgt=self.tds.text\_to\_i(s[1])131neighbors=torch.stack([torch.stack([self.tds.text\_to\_i(n)forninchunks])forchunksins[2]])

133returnsrc,tgt,neighbors

136if\_\_name\_\_=='\_\_main\_\_':137build\_dataset()