Select random query vectors from the dataset
(vectors, num_queries)
| 170 | |
| 171 | |
| 172 | def select_query_vectors(vectors, num_queries): |
| 173 | """Select random query vectors from the dataset""" |
| 174 | if os.path.exists(QUERY_VECTORS_FILE): |
| 175 | print(f"Loading existing query vectors from {QUERY_VECTORS_FILE}") |
| 176 | with open(QUERY_VECTORS_FILE, "rb") as f: |
| 177 | query_indices = pickle.load(f) |
| 178 | return [vectors[i] for i in query_indices] |
| 179 | |
| 180 | print(f"Selecting {num_queries} random query vectors...") |
| 181 | query_indices = random.sample(range(len(vectors)), num_queries) |
| 182 | query_vectors = [vectors[i] for i in query_indices] |
| 183 | |
| 184 | # Save indices to disk |
| 185 | with open(QUERY_VECTORS_FILE, "wb") as f: |
| 186 | pickle.dump(query_indices, f) |
| 187 | |
| 188 | print(f"Query vectors selected and indices saved to {QUERY_VECTORS_FILE}") |
| 189 | return query_vectors |
| 190 | |
| 191 | |
| 192 | def compute_brute_force_results(vectors, query_vectors, dimension, top_k=10): |