Build a FAISS index from an Arrow file containing keys and values. Args: dstore_path: Path to the Arrow file containing the datastore num_keys_to_add_at_a_time: Number of keys to add at a time ncentroids: Number of centroids for IVFPQ seed: Random seed
(
dstore_path,
num_keys_to_add_at_a_time=1_000_000,
ncentroids=4096,
seed=42,
code_size=32,
probe=8
)
| 124 | return np.vstack(all_samples) |
| 125 | |
| 126 | def build_index( |
| 127 | dstore_path, |
| 128 | num_keys_to_add_at_a_time=1_000_000, |
| 129 | ncentroids=4096, |
| 130 | seed=42, |
| 131 | code_size=32, |
| 132 | probe=8 |
| 133 | ): |
| 134 | """ |
| 135 | Build a FAISS index from an Arrow file containing keys and values. |
| 136 | |
| 137 | Args: |
| 138 | dstore_path: Path to the Arrow file containing the datastore |
| 139 | num_keys_to_add_at_a_time: Number of keys to add at a time |
| 140 | ncentroids: Number of centroids for IVFPQ |
| 141 | seed: Random seed |
| 142 | code_size: Code size for PQ |
| 143 | probe: Number of probes for query |
| 144 | """ |
| 145 | # Parse the dstore path to extract necessary information |
| 146 | dstore_info = parse_dstore_path(dstore_path) |
| 147 | logger.info(f"Parsed dstore path: {dstore_info}") |
| 148 | |
| 149 | dimension = dstore_info["dimension"] |
| 150 | |
| 151 | logger.info('Loading Dataset...') |
| 152 | dstore = Dataset.from_file(dstore_path) |
| 153 | # Set format to numpy for proper array conversion |
| 154 | dstore.set_format(type='numpy', columns=['keys', 'vals']) |
| 155 | |
| 156 | # Save dstore.vals to a separate pickle file for later use |
| 157 | vals_path = os.path.join(dstore_info["dstore_dir"], f"{dstore_info['eval_subset']}_vals.pkl") |
| 158 | |
| 159 | # Select only vals column |
| 160 | vals_dataset = dstore.select_columns(['vals']) |
| 161 | vals_dataset.set_format(type='torch', columns=['vals']) |
| 162 | |
| 163 | # Note that since datasets version 4.0.0, we can't use direct column selecting since the implementation of lazy columns, see pr https://github.com/huggingface/datasets/pull/7614 |
| 164 | vals_tensor = vals_dataset[:]['vals'] |
| 165 | logger.info(f"Saved val tensor shape: {vals_tensor.shape}") |
| 166 | with open(vals_path, 'wb') as f: |
| 167 | pickle.dump(vals_tensor, f) |
| 168 | |
| 169 | logger.info('Building index...') |
| 170 | index_name = get_index_path(dstore_info) |
| 171 | logger.info(f"Index will be saved to: {index_name}") |
| 172 | |
| 173 | max_threads = multiprocessing.cpu_count() |
| 174 | faiss.omp_set_num_threads(max_threads) |
| 175 | logger.info(f'Total CPU count: {max_threads}') |
| 176 | |
| 177 | quantizer = faiss.IndexFlatL2(dimension) |
| 178 | index = faiss.IndexIVFPQ(quantizer, dimension, ncentroids, code_size, 8) |
| 179 | index.nprobe = probe |
| 180 | |
| 181 | logger.info('Training Index...') |
| 182 | sample_size = min(1000000, len(dstore)) |
| 183 | train_data = select_continuous_chunks(dstore, total_sample_size=sample_size, num_chunks=1000, seed=seed) |
no test coverage detected