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