16fromtypingimportList,Optional1718importfaiss19importnumpyasnp20importtorch2122fromlabmlimportlab,monit23fromlabml\_helpers.datasets.textimportTextFileDataset24fromlabml\_nn.transformers.retro.bert\_embeddingsimportBERTChunkEmbeddings

27defbuild\_database(chunk\_len:int=16,batch\_size:int=64,d\_emb:int=768,n\_centeroids:int=256,28code\_size:int=64,n\_probe:int=8,n\_train:int=50\_000):

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

49text=dataset.train

52chunks=[text[i:i+chunk\_len]foriinrange(0,len(text),chunk\_len)ifi+chunk\_len\*2\<len(text)]

54chunk\_offsets=np.array([iforiinrange(0,len(text),chunk\_len)ifi+chunk\_len\*2\<len(text)])

56n\_chunks=len(chunks)

59bert=BERTChunkEmbeddings(torch.device('cuda:0'))

62chunk\_emb=[]63foriinmonit.iterate('Get embeddings',range(0,n\_chunks,batch\_size)):64chunk\_emb.append(bert(chunks[i:i+batch\_size]).cpu())

66chunk\_emb=torch.cat(chunk\_emb,dim=0).numpy()

69quantizer=faiss.IndexFlatL2(d\_emb)70index=faiss.IndexIVFPQ(quantizer,d\_emb,n\_centeroids,code\_size,8)71index.nprobe=n\_probe

74random\_sample=np.random.choice(np.arange(n\_chunks),size=[min(n\_train,n\_chunks)],replace=False)

77withmonit.section('Train index'):78index.train(chunk\_emb[random\_sample])

81forsinmonit.iterate('Index',range(0,n\_chunks,1024)):82e=min(s+1024,n\_chunks)

84index.add\_with\_ids(chunk\_emb[s:e],chunk\_offsets[s:e])

87withmonit.section('Save'):88faiss.write\_index(index,str(lab.get\_data\_path()/'retro.index'))

91classRetroIndex:

96def\_\_init\_\_(self,chunk\_len:int=16,n\_probe:int=8,97n\_neighbors:int=2,n\_extra:int=2,98exclude\_neighbor\_span:int=8):

108self.n\_neighbors=n\_neighbors109self.chunk\_len=chunk\_len110self.exclude\_neighbor\_span=exclude\_neighbor\_span111self.n\_extra=n\_extra

114self.bert=BERTChunkEmbeddings(torch.device('cuda:0'))

116withmonit.section('Load index'):117self.index=faiss.read\_index(str(lab.get\_data\_path()/'retro.index'))118self.index.nprobe=n\_probe

120deffilter\_neighbors(self,offset:int,neighbor\_offsets:List[int]):

127return[nforninneighbor\_offsets128ifn\<offset-(self.chunk\_len+self.exclude\_neighbor\_span)129orn\>offset+(self.chunk\_len+self.exclude\_neighbor\_span)]

131def\_\_call\_\_(self,query\_chunks:List[str],offsets:Optional[List[int]]):

137emb=self.bert(query\_chunks).cpu()

140distance,neighbor\_offsets=self.index.search(emb.numpy(),self.n\_neighbors+self.n\_extra)

143ifoffsetsisnotNone:144neighbor\_offsets=[self.filter\_neighbors(off,n\_off)145foroff,n\_offinzip(offsets,neighbor\_offsets)]

148neighbor\_offsets=[n\_off[:self.n\_neighbors]forn\_offinneighbor\_offsets]

151returnneighbor\_offsets

155if\_\_name\_\_=='\_\_main\_\_':156build\_database()