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