Add distribution to all embedders

milli/src/vector/manual.rs

@@ -1,19 +1,21 @@
 use super::error::EmbedError;
-use super::Embeddings;
+use super::{DistributionShift, Embeddings};
 
 #[derive(Debug, Clone, Copy)]
 pub struct Embedder {
     dimensions: usize,
+    distribution: Option<DistributionShift>,
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
 pub struct EmbedderOptions {
     pub dimensions: usize,
+    pub distribution: Option<DistributionShift>,
 }
 
 impl Embedder {
     pub fn new(options: EmbedderOptions) -> Self {
-        Self { dimensions: options.dimensions }
+        Self { dimensions: options.dimensions, distribution: options.distribution }
     }
 
     pub fn embed(&self, mut texts: Vec<String>) -> Result<Vec<Embeddings<f32>>, EmbedError> {
@@ -31,4 +33,8 @@ impl Embedder {
     ) -> Result<Vec<Vec<Embeddings<f32>>>, EmbedError> {
         text_chunks.into_iter().map(|prompts| self.embed(prompts)).collect()
     }
+
+    pub fn distribution(&self) -> Option<DistributionShift> {
+        self.distribution
+    }
 }