| 12 | using gt_type = rabitqlib::RowMajorArray<uint32_t>; |
| 13 | |
| 14 | int main(int argc, char** argv) { |
| 15 | if (argc < 6) { |
| 16 | std::cerr << "Usage: " << argv[0] << " <arg1> <arg2> <arg3> <arg4> <arg5>\n" |
| 17 | << "arg1: path for data file, format .fvecs\n" |
| 18 | << "arg2: path for centroids file generated by ivf.py\n" |
| 19 | << "arg3: path for cluster ids file generated by ivf.py\n" |
| 20 | << "arg4: total number of bits for quantization\n" |
| 21 | << "arg5: path for saving index\n" |
| 22 | << "arg6: metric type (\"l2\" or \"ip\"), l2 by default\n" |
| 23 | << "arg7: if use faster quantization (\"true\" or \"false\"), false by " |
| 24 | "default\n"; |
| 25 | exit(1); |
| 26 | } |
| 27 | |
| 28 | rabitqlib::MetricType metric_type = rabitqlib::METRIC_L2; |
| 29 | if (argc > 6) { |
| 30 | std::string metric_str(argv[6]); |
| 31 | if (metric_str == "ip" || metric_str == "IP") { |
| 32 | metric_type = rabitqlib::METRIC_IP; |
| 33 | } |
| 34 | } |
| 35 | if (metric_type == rabitqlib::METRIC_IP) { |
| 36 | std::cout << "Metric Type: IP\n"; |
| 37 | } else if (metric_type == rabitqlib::METRIC_L2) { |
| 38 | std::cout << "Metric Type: L2\n"; |
| 39 | } |
| 40 | |
| 41 | bool faster_quant = false; |
| 42 | if (argc > 7) { |
| 43 | std::string faster_str(argv[7]); |
| 44 | if (faster_str == "true") { |
| 45 | faster_quant = true; |
| 46 | std::cout << "Using faster quantize for indexing...\n"; |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | char* data_file = argv[1]; |
| 51 | char* centroids_file = argv[2]; |
| 52 | char* cids_file = argv[3]; |
| 53 | size_t total_bits = atoi(argv[4]); |
| 54 | char* index_file = argv[5]; |
| 55 | |
| 56 | data_type data; |
| 57 | data_type centroids; |
| 58 | gt_type cids; |
| 59 | |
| 60 | rabitqlib::load_vecs<float, data_type>(data_file, data); |
| 61 | rabitqlib::load_vecs<float, data_type>(centroids_file, centroids); |
| 62 | rabitqlib::load_vecs<PID, gt_type>(cids_file, cids); |
| 63 | |
| 64 | size_t num_points = data.rows(); |
| 65 | size_t dim = data.cols(); |
| 66 | size_t k = centroids.rows(); |
| 67 | |
| 68 | std::cout << "data loaded\n"; |
| 69 | std::cout << "\tN: " << num_points << '\n'; |
| 70 | std::cout << "\tDIM: " << dim << '\n'; |
| 71 | |