Generate dataset if not already on disk
(num_vectors, dimension, max_non_zero_dims)
| 157 | |
| 158 | |
| 159 | def generate_dataset(num_vectors, dimension, max_non_zero_dims): |
| 160 | """Generate dataset if not already on disk""" |
| 161 | if os.path.exists(DATASET_FILE): |
| 162 | print(f"Loading existing dataset from {DATASET_FILE}") |
| 163 | with open(DATASET_FILE, "rb") as f: |
| 164 | vectors = pickle.load(f) |
| 165 | return vectors |
| 166 | |
| 167 | print(f"Generating {num_vectors} random sparse vectors...") |
| 168 | vectors = [ |
| 169 | generate_random_sparse_vector(id, dimension, max_non_zero_dims) |
| 170 | for id in tqdm(range(num_vectors)) |
| 171 | ] |
| 172 | |
| 173 | # Save to disk |
| 174 | with open(DATASET_FILE, "wb") as f: |
| 175 | pickle.dump(vectors, f) |
| 176 | |
| 177 | print(f"Dataset generated and saved to {DATASET_FILE}") |
| 178 | return vectors |
| 179 | |
| 180 | |
| 181 | def select_query_vectors(vectors, num_queries): |
no test coverage detected