| 12 | using gt_type = rabitqlib::RowMajorArray<uint32_t>; |
| 13 | |
| 14 | int main(int argc, char** argv) { |
| 15 | if (argc < 5) { |
| 16 | std::cerr << "Usage: " << argv[0] << " <arg1> <arg2> <arg3> <arg4>\n" |
| 17 | << "arg1: path for data file, format .fvecs\n" |
| 18 | << "arg2: degree bound for symqg, must be a multiple of 32\n" |
| 19 | << "arg3: ef for indexing \n" |
| 20 | << "arg4: path for saving index\n" |
| 21 | << "arg5: metric type (\"l2\" or \"ip\"), l2 by default\n"; |
| 22 | exit(1); |
| 23 | } |
| 24 | |
| 25 | char* data_file = argv[1]; |
| 26 | size_t degree = atoi(argv[2]); |
| 27 | size_t ef = atoi(argv[3]); |
| 28 | char* index_file = argv[4]; |
| 29 | |
| 30 | rabitqlib::MetricType metric_type = rabitqlib::METRIC_L2; |
| 31 | if (argc > 5) { |
| 32 | std::string metric_str(argv[5]); |
| 33 | if (metric_str == "ip" || metric_str == "IP") { |
| 34 | metric_type = rabitqlib::METRIC_IP; |
| 35 | } |
| 36 | } |
| 37 | if (metric_type == rabitqlib::METRIC_IP) { |
| 38 | std::cout << "Metric Type: IP\n"; |
| 39 | } else if (metric_type == rabitqlib::METRIC_L2) { |
| 40 | std::cout << "Metric Type: L2\n"; |
| 41 | } |
| 42 | |
| 43 | data_type data; |
| 44 | |
| 45 | rabitqlib::load_vecs<float, data_type>(data_file, data); |
| 46 | |
| 47 | rabitqlib::StopW stopw; |
| 48 | |
| 49 | index_type qg(data.rows(), data.cols(), degree, metric_type); |
| 50 | |
| 51 | rabitqlib::symqg::QGBuilder builder(qg, ef, data.data()); |
| 52 | |
| 53 | // 3 iters, refine at last iter |
| 54 | builder.build(); |
| 55 | |
| 56 | auto milisecs = stopw.get_elapsed_mili(); |
| 57 | |
| 58 | std::cout << "Indexing time " << milisecs / 1000.F << " secs\n"; |
| 59 | |
| 60 | qg.save(index_file); |
| 61 | |
| 62 | return 0; |
| 63 | } |