(args=None)
| 16 | |
| 17 | |
| 18 | def main(args=None) -> None: |
| 19 | # 1. Load queries and ground truth |
| 20 | queries = read_fvecs(args.query_file) |
| 21 | gt = read_ivecs(args.gt_file) |
| 22 | nq = queries.shape[0] |
| 23 | print(f"Queries: {queries.shape}, GT: {gt.shape}") |
| 24 | print(f"TopK: {args.topk}, use_hacc: {args.use_hacc}") |
| 25 | |
| 26 | # 2. Load index |
| 27 | idx = IvfIndex.load(args.index_file) |
| 28 | print(f"Index loaded — dim={idx.dim}, clusters={idx.num_clusters}") |
| 29 | |
| 30 | all_qps = np.zeros((args.test_rounds, len(NPROBES))) |
| 31 | all_recall = np.zeros((args.test_rounds, len(NPROBES))) |
| 32 | |
| 33 | # 3. Benchmark |
| 34 | for r in range(args.test_rounds): |
| 35 | for l, nprobe in enumerate(NPROBES): |
| 36 | if nprobe > idx.num_clusters: |
| 37 | print(f"nprobe {nprobe} is larger than number of clusters, " |
| 38 | f"will use nprobe = num_clusters ({idx.num_clusters}).") |
| 39 | |
| 40 | t0 = time() |
| 41 | ids, _ = idx.search(queries, k=args.topk, nprobe=nprobe, high_accuracy=args.use_hacc, num_threads=args.num_threads) |
| 42 | elapsed = time() - t0 |
| 43 | |
| 44 | all_qps[r, l] = nq / elapsed |
| 45 | all_recall[r, l] = compute_recall(ids, gt, args.topk) |
| 46 | |
| 47 | avg_qps = all_qps.mean(axis=0) |
| 48 | avg_recall = all_recall.mean(axis=0) |
| 49 | |
| 50 | # 4. Print results table |
| 51 | print(f"\n{'nprobe':<10}{'QPS':<14}{'Recall'}") |
| 52 | print("-" * 35) |
| 53 | for i, nprobe in enumerate(NPROBES): |
| 54 | print(f"{nprobe:<10}{avg_qps[i]:<14.1f}{avg_recall[i]:.4f}") |
| 55 | |
| 56 | |
| 57 | if __name__ == "__main__": |
no test coverage detected