Function main

sample/python/hnsw_rabitq_querying.py:17–53

(args=None)

def main(args=None) -> None:
    # 1. Load queries and ground truth
    queries = read_fvecs(args.query_file)
    gt      = read_ivecs(args.gt_file)
    nq      = queries.shape[0]
    print(f"Queries: {queries.shape}, GT: {gt.shape}")

    # 2. Load index
    idx = HnswIndex.load(args.index_file)
    print(f"Index loaded — dim={idx.dim}")
    print(f"TopK: {args.topk}")

    print("\nsearch start >.....\n")

    all_qps    = np.zeros((args.test_rounds, len(EFS)))
    all_recall = np.zeros((args.test_rounds, len(EFS)))

    for i_probe, ef in enumerate(EFS):
        for r in range(args.test_rounds):
            t0 = time()
            ids, _ = idx.search(queries, k=args.topk, ef=ef, num_threads=args.num_threads)
            elapsed = time() - t0  # seconds

            qps    = nq / elapsed
            recall = compute_recall(ids, gt, args.topk)

            all_qps[r, i_probe]    = qps
            all_recall[r, i_probe] = recall

    avg_qps    = all_qps.mean(axis=0)
    avg_recall = all_recall.mean(axis=0)

    # 3. Print results table
    print(f"{'EF':<8}{'QPS':<14}{'Recall'}")
    print("-" * 35)
    for i, ef in enumerate(EFS):
        print(f"{ef:<8}{avg_qps[i]:<14.1f}{avg_recall[i]:<12.4f}")


if __name__ == "__main__":