| 17 | using gt_type = rabitqlib::RowMajorArray<uint32_t>; |
| 18 | |
| 19 | int main(int argc, char* argv[]) { |
| 20 | if (argc < 3) { |
| 21 | std::cerr << "Usage: " << argv[0] << " <arg1> <arg2> <arg3>\n" |
| 22 | << "arg1: path for index \n" |
| 23 | << "arg2: path for query file, format .fvecs\n" |
| 24 | << "arg3: path for groundtruth file format .ivecs\n"; |
| 25 | exit(1); |
| 26 | } |
| 27 | |
| 28 | char* index_file = argv[1]; |
| 29 | char* query_file = argv[2]; |
| 30 | char* gt_file = argv[3]; |
| 31 | |
| 32 | data_type query; |
| 33 | gt_type gt; |
| 34 | rabitqlib::load_vecs<float, data_type>(query_file, query); |
| 35 | rabitqlib::load_vecs<uint32_t, gt_type>(gt_file, gt); |
| 36 | size_t nq = query.rows(); |
| 37 | size_t total_count = nq * topk; |
| 38 | |
| 39 | index_type hnsw; |
| 40 | |
| 41 | hnsw.load(index_file); |
| 42 | |
| 43 | rabitqlib::StopW stopw; |
| 44 | |
| 45 | auto nefs = efs; |
| 46 | |
| 47 | size_t length = nefs.size(); |
| 48 | |
| 49 | std::vector<std::vector<float>> all_qps(test_round, std::vector<float>(length)); |
| 50 | std::vector<std::vector<float>> all_recall(test_round, std::vector<float>(length)); |
| 51 | |
| 52 | std::cout << "search start >.....\n"; |
| 53 | |
| 54 | for (size_t i_probe = 0; i_probe < length; ++i_probe) { |
| 55 | for (size_t r = 0; r < test_round; r++) { |
| 56 | size_t ef = nefs[i_probe]; |
| 57 | size_t total_correct = 0; |
| 58 | float total_time = 0; |
| 59 | |
| 60 | auto start = std::chrono::high_resolution_clock::now(); |
| 61 | |
| 62 | std::vector<std::vector<std::pair<float, PID>>> res = |
| 63 | hnsw.search(query.data(), nq, topk, ef, 1); |
| 64 | |
| 65 | auto end = std::chrono::high_resolution_clock::now(); |
| 66 | |
| 67 | float elapsed_us = |
| 68 | std::chrono::duration<float, std::micro>(end - start).count(); |
| 69 | |
| 70 | total_time += elapsed_us; |
| 71 | |
| 72 | for (size_t i = 0; i < nq; i++) { |
| 73 | for (size_t j = 0; j < topk; j++) { |
| 74 | for (size_t k = 0; k < topk; k++) { |
| 75 | if (gt(i, k) == res[i][j].second) { |
| 76 | total_correct++; |