15fromtypingimportOptional1617importfaiss18importnumpyasnp19importtorch2021fromlabmlimportexperiment,monit,lab22fromlabml.utils.pytorchimportget\_modules23fromlabml\_nn.transformers.knn.train\_modelimportConfigs

26defload\_experiment(run\_uuid:str,checkpoint:Optional[int]=None):

32conf=Configs()

34conf\_dict=experiment.load\_configs(run\_uuid)

36conf\_dict['is\_save\_ff\_input']=True

39experiment.evaluate()

41experiment.configs(conf,conf\_dict)

43experiment.add\_pytorch\_models(get\_modules(conf))

45experiment.load(run\_uuid,checkpoint)

48experiment.start()4950returnconf

53defgather\_keys(conf:Configs):

62d\_model=conf.transformer.d\_model

64data\_loader=conf.trainer.data\_loader

67n\_keys=data\_loader.data.shape[0]\*data\_loader.data.shape[1]-1

69keys\_store=np.memmap(str(lab.get\_data\_path()/'keys.npy'),dtype=np.float32,mode='w+',shape=(n\_keys,d\_model))

71vals\_store=np.memmap(str(lab.get\_data\_path()/'vals.npy'),dtype=np.int,mode='w+',shape=(n\_keys,1))

74added=075withtorch.no\_grad():

77fori,batchinmonit.enum("Collect data",data\_loader,is\_children\_silent=True):

79vals=batch[1].view(-1,1)

81data=batch[0].to(conf.device)

83\_=conf.model(data)

85keys=conf.model.ff\_input.view(-1,d\_model)

87keys\_store[added:added+keys.shape[0]]=keys.cpu()

89vals\_store[added:added+keys.shape[0]]=vals

91added+=keys.shape[0]

94defbuild\_index(conf:Configs,n\_centeroids:int=2048,code\_size:int=64,n\_probe:int=8,n\_train:int=200\_000):

104d\_model=conf.transformer.d\_model

106data\_loader=conf.trainer.data\_loader

109n\_keys=data\_loader.data.shape[0]\*data\_loader.data.shape[1]-1

113quantizer=faiss.IndexFlatL2(d\_model)114index=faiss.IndexIVFPQ(quantizer,d\_model,n\_centeroids,code\_size,8)115index.nprobe=n\_probe

118keys\_store=np.memmap(str(lab.get\_data\_path()/'keys.npy'),dtype=np.float32,mode='r',shape=(n\_keys,d\_model))

121random\_sample=np.random.choice(np.arange(n\_keys),size=[min(n\_train,n\_keys)],replace=False)122123withmonit.section('Train index'):

125index.train(keys\_store[random\_sample])

128forsinmonit.iterate('Index',range(0,n\_keys,1024)):129e=min(s+1024,n\_keys)

131keys=keys\_store[s:e]

133idx=np.arange(s,e)

135index.add\_with\_ids(keys,idx)136137withmonit.section('Save'):

139faiss.write\_index(index,str(lab.get\_data\_path()/'faiss.index'))

142defmain():

145conf=load\_experiment('4984b85c20bf11eb877a69c1a03717cd')

147conf.model.eval()

150gather\_keys(conf)

152build\_index(conf)153154155if\_\_name\_\_=='\_\_main\_\_':156main()