Method build_index

knn_utils/saveEmbedMulti.py:440–477

(self, num_keys_to_add_at_a_time=1_000_000, ncentroids=4096, seed=42, code_size=32, probe=8)

            self.model.broken_into = None
        
    def build_index(self, num_keys_to_add_at_a_time=1_000_000, ncentroids=4096, seed=42, code_size=32, probe=8):
        logger.info('Loading Dataset...')
        dstore_path = self._get_arrow_file_path()
        dstore = Dataset.from_file(dstore_path)
        # Set format to numpy for proper array conversion
        dstore.set_format(type='numpy', columns=['keys', 'vals'])
        
        logger.info('Building index...')
        index_name = get_index_path(self.dstore_dir, self.model.config.model_type, self.eval_subset, self.dimension)
        
        quantizer = faiss.IndexFlatL2(self.dimension)
        index = faiss.IndexIVFPQ(quantizer, self.dimension, ncentroids, code_size, 8)
        index.nprobe = probe
        
        logger.info('Training Index...')
        np.random.seed(seed)
        sample_size = min(200000, len(dstore))
        random_sample = np.random.choice(len(dstore), size=sample_size, replace=False)
        
        train_data = np.array(dstore.select(random_sample)['keys']).astype(np.float32)
        start = time.time()
        index.train(train_data)
        logger.info(f'Training took {time.time() - start:.2f} s')
        
        logger.info('Adding Keys...')
        start_time = time.time()
        
        for start in tqdm(range(0, len(dstore), num_keys_to_add_at_a_time)):
            end = min(len(dstore), start + num_keys_to_add_at_a_time)
            to_add = np.array(dstore.select(range(start,end))['keys']).astype(np.float32)  
            
            index.add_with_ids(to_add, np.arange(start, end))
            
            if (start // num_keys_to_add_at_a_time) % 10 == 0:
                faiss.write_index(index, index_name)
        
        faiss.write_index(index, index_name)
        logger.info(f'Added {len(dstore)} keys in {time.time() - start_time:.2f} s')

class ActivationCapturer(nn.Module):
    def __init__(self, layer, capture_input=False):