| 36 | )) {} |
| 37 | |
| 38 | void build( |
| 39 | py::handle data, |
| 40 | py::handle centroids, |
| 41 | py::handle cluster_ids, |
| 42 | size_t num_threads = 1, |
| 43 | bool fast_quantization = false |
| 44 | ) { |
| 45 | auto data_array = ensure_2d_array<float>(data, "data"); |
| 46 | auto centroids_array = ensure_2d_array<float>(centroids, "centroids"); |
| 47 | auto cluster_ids_array = ensure_1d_array<rabitqlib::PID>(cluster_ids, "cluster_ids"); |
| 48 | |
| 49 | if (static_cast<size_t>(data_array.shape(1)) != dim_) { |
| 50 | throw std::invalid_argument("data dimension does not match index dim"); |
| 51 | } |
| 52 | if (static_cast<size_t>(centroids_array.shape(1)) != dim_) { |
| 53 | throw std::invalid_argument("centroid dimension does not match index dim"); |
| 54 | } |
| 55 | if (static_cast<size_t>(cluster_ids_array.shape(0)) != static_cast<size_t>(data_array.shape(0))) { |
| 56 | throw std::invalid_argument("cluster_ids length must match number of rows in data"); |
| 57 | } |
| 58 | |
| 59 | index_->construct(data_array.data(), centroids_array.data(), cluster_ids_array.data(), fast_quantization, num_threads); |
| 60 | built_ = true; |
| 61 | } |
| 62 | |
| 63 | py::tuple search( |
| 64 | py::handle queries, |