Train the index from a set of vectors.
(&mut self, vectors: &[&[f32]])
| 65 | |
| 66 | /// Train the index from a set of vectors. |
| 67 | pub fn train(&mut self, vectors: &[&[f32]]) { |
| 68 | assert!(!vectors.is_empty()); |
| 69 | assert!(self.dim > 0); |
| 70 | assert!( |
| 71 | self.dim.is_multiple_of(self.params.pq_m), |
| 72 | "dim {} must be divisible by pq_m {}", |
| 73 | self.dim, |
| 74 | self.params.pq_m |
| 75 | ); |
| 76 | |
| 77 | let n_cells = self.params.n_cells.min(vectors.len()); |
| 78 | self.centroids = kmeans_centroids(vectors, self.dim, n_cells, 20); |
| 79 | self.cells = vec![Vec::new(); self.centroids.len()]; |
| 80 | |
| 81 | let mut residuals: Vec<Vec<f32>> = Vec::with_capacity(vectors.len()); |
| 82 | for v in vectors { |
| 83 | let cell = self.nearest_centroid(v); |
| 84 | let res: Vec<f32> = v |
| 85 | .iter() |
| 86 | .zip(&self.centroids[cell]) |
| 87 | .map(|(a, b)| a - b) |
| 88 | .collect(); |
| 89 | residuals.push(res); |
| 90 | } |
| 91 | let res_refs: Vec<&[f32]> = residuals.iter().map(|r| r.as_slice()).collect(); |
| 92 | self.pq = Some(PqCodec::train( |
| 93 | &res_refs, |
| 94 | self.dim, |
| 95 | self.params.pq_m, |
| 96 | self.params.pq_k, |
| 97 | 20, |
| 98 | )); |
| 99 | } |
| 100 | |
| 101 | /// Add a vector to the index. Returns the assigned ID. |
| 102 | pub fn add(&mut self, vector: &[f32]) -> u32 { |