Train PQ codebooks from a set of training vectors via k-means. `m` = number of subvectors (must divide `dim` evenly). `k` = centroids per subvector (typically 256). `max_iter` = k-means iterations (20 is usually sufficient).
(vectors: &[&[f32]], dim: usize, m: usize, k: usize, max_iter: usize)
| 82 | /// `k` = centroids per subvector (typically 256). |
| 83 | /// `max_iter` = k-means iterations (20 is usually sufficient). |
| 84 | pub fn train(vectors: &[&[f32]], dim: usize, m: usize, k: usize, max_iter: usize) -> Self { |
| 85 | assert!(!vectors.is_empty()); |
| 86 | assert!(dim > 0 && m > 0 && k > 0); |
| 87 | assert!( |
| 88 | dim.is_multiple_of(m), |
| 89 | "dim ({dim}) must be divisible by m ({m})" |
| 90 | ); |
| 91 | |
| 92 | let sub_dim = dim / m; |
| 93 | let mut codebooks = Vec::with_capacity(m); |
| 94 | |
| 95 | for sub in 0..m { |
| 96 | let offset = sub * sub_dim; |
| 97 | // Extract sub-vectors for this subspace. |
| 98 | let sub_vectors: Vec<&[f32]> = vectors |
| 99 | .iter() |
| 100 | .map(|v| &v[offset..offset + sub_dim]) |
| 101 | .collect(); |
| 102 | |
| 103 | let centroids = kmeans(&sub_vectors, sub_dim, k, max_iter); |
| 104 | codebooks.push(centroids); |
| 105 | } |
| 106 | |
| 107 | Self { |
| 108 | dim, |
| 109 | m, |
| 110 | k, |
| 111 | sub_dim, |
| 112 | codebooks, |
| 113 | governor: None, |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | /// Encode a vector: for each subvector, find the nearest centroid index. |
| 118 | /// |
no test coverage detected