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