Merge #5355

5355: Support fetching the pooling method from the model configuration r=Kerollmops a=dureuill

# Pull Request

## Related issue
Fixes #5354 

## What does this PR do?
- Fetches the pooling configuration from the model repository
- Use a pooling method that depends on the pooling configuration of that model.
- Allow overriding the pooling method with a new huggingFace embedder parameter `pooling`
  - for backward-compatibility with Meilisearch v1.13
  - for compatibility with embedders that exhibit the same behavior as Meilisearch v1.13
- Handle the default value of that new parameter
   - for compatibility, when importing a db/a dump, it should be set to `forceMean`
   - when (re)set from the settings for an embedder, it should be set to `useModel`


Co-authored-by: Louis Dureuil <louis@meilisearch.com>

crates/milli/src/update/index_documents/mod.rs

@@ -2768,6 +2768,7 @@ mod tests {
                         source: Setting::Set(crate::vector::settings::EmbedderSource::UserProvided),
                         model: Setting::NotSet,
                         revision: Setting::NotSet,
+                        pooling: Setting::NotSet,
                         api_key: Setting::NotSet,
                         dimensions: Setting::Set(3),
                         document_template: Setting::NotSet,

crates/milli/src/update/settings.rs

@@ -1676,6 +1676,7 @@ fn validate_prompt(
             source,
             model,
             revision,
+            pooling,
             api_key,
             dimensions,
             document_template: Setting::Set(template),
@@ -1709,6 +1710,7 @@ fn validate_prompt(
                 source,
                 model,
                 revision,
+                pooling,
                 api_key,
                 dimensions,
                 document_template: Setting::Set(template),
@@ -1735,6 +1737,7 @@ pub fn validate_embedding_settings(
         source,
         model,
         revision,
+        pooling,
         api_key,
         dimensions,
         document_template,
@@ -1776,6 +1779,7 @@ pub fn validate_embedding_settings(
             source,
             model,
             revision,
+            pooling,
             api_key,
             dimensions,
             document_template,
@@ -1791,6 +1795,7 @@ pub fn validate_embedding_settings(
     match inferred_source {
         EmbedderSource::OpenAi => {
             check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?;
+            check_unset(&pooling, EmbeddingSettings::POOLING, inferred_source, name)?;
 
             check_unset(&request, EmbeddingSettings::REQUEST, inferred_source, name)?;
             check_unset(&response, EmbeddingSettings::RESPONSE, inferred_source, name)?;
@@ -1829,6 +1834,7 @@ pub fn validate_embedding_settings(
         EmbedderSource::Ollama => {
             check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?;
             check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?;
+            check_unset(&pooling, EmbeddingSettings::POOLING, inferred_source, name)?;
 
             check_unset(&request, EmbeddingSettings::REQUEST, inferred_source, name)?;
             check_unset(&response, EmbeddingSettings::RESPONSE, inferred_source, name)?;
@@ -1846,6 +1852,7 @@ pub fn validate_embedding_settings(
         EmbedderSource::UserProvided => {
             check_unset(&model, EmbeddingSettings::MODEL, inferred_source, name)?;
             check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?;
+            check_unset(&pooling, EmbeddingSettings::POOLING, inferred_source, name)?;
             check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?;
             check_unset(
                 &document_template,
@@ -1869,6 +1876,7 @@ pub fn validate_embedding_settings(
         EmbedderSource::Rest => {
             check_unset(&model, EmbeddingSettings::MODEL, inferred_source, name)?;
             check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?;
+            check_unset(&pooling, EmbeddingSettings::POOLING, inferred_source, name)?;
             check_set(&url, EmbeddingSettings::URL, inferred_source, name)?;
             check_set(&request, EmbeddingSettings::REQUEST, inferred_source, name)?;
             check_set(&response, EmbeddingSettings::RESPONSE, inferred_source, name)?;
@@ -1878,6 +1886,7 @@ pub fn validate_embedding_settings(
         source,
         model,
         revision,
+        pooling,
         api_key,
         dimensions,
         document_template,

crates/milli/src/vector/error.rs

@@ -262,6 +262,31 @@ impl NewEmbedderError {
         }
     }
 
+    pub fn open_pooling_config(
+        pooling_config_filename: PathBuf,
+        inner: std::io::Error,
+    ) -> NewEmbedderError {
+        let open_config = OpenPoolingConfig { filename: pooling_config_filename, inner };
+
+        Self {
+            kind: NewEmbedderErrorKind::OpenPoolingConfig(open_config),
+            fault: FaultSource::Runtime,
+        }
+    }
+
+    pub fn deserialize_pooling_config(
+        model_name: String,
+        pooling_config_filename: PathBuf,
+        inner: serde_json::Error,
+    ) -> NewEmbedderError {
+        let deserialize_pooling_config =
+            DeserializePoolingConfig { model_name, filename: pooling_config_filename, inner };
+        Self {
+            kind: NewEmbedderErrorKind::DeserializePoolingConfig(deserialize_pooling_config),
+            fault: FaultSource::Runtime,
+        }
+    }
+
     pub fn open_tokenizer(
         tokenizer_filename: PathBuf,
         inner: Box<dyn std::error::Error + Send + Sync>,
@@ -319,6 +344,13 @@ pub struct OpenConfig {
     pub inner: std::io::Error,
 }
 
+#[derive(Debug, thiserror::Error)]
+#[error("could not open pooling config at {filename}: {inner}")]
+pub struct OpenPoolingConfig {
+    pub filename: PathBuf,
+    pub inner: std::io::Error,
+}
+
 #[derive(Debug, thiserror::Error)]
 #[error("for model '{model_name}', could not deserialize config at {filename} as JSON: {inner}")]
 pub struct DeserializeConfig {
@@ -327,6 +359,14 @@ pub struct DeserializeConfig {
     pub inner: serde_json::Error,
 }
 
+#[derive(Debug, thiserror::Error)]
+#[error("for model '{model_name}', could not deserialize file at `{filename}` as a pooling config: {inner}")]
+pub struct DeserializePoolingConfig {
+    pub model_name: String,
+    pub filename: PathBuf,
+    pub inner: serde_json::Error,
+}
+
 #[derive(Debug, thiserror::Error)]
 #[error("model `{model_name}` appears to be unsupported{}\n  - inner error: {inner}",
 if architectures.is_empty() {
@@ -354,8 +394,12 @@ pub enum NewEmbedderErrorKind {
     #[error(transparent)]
     OpenConfig(OpenConfig),
     #[error(transparent)]
+    OpenPoolingConfig(OpenPoolingConfig),
+    #[error(transparent)]
     DeserializeConfig(DeserializeConfig),
     #[error(transparent)]
+    DeserializePoolingConfig(DeserializePoolingConfig),
+    #[error(transparent)]
     UnsupportedModel(UnsupportedModel),
     #[error(transparent)]
     OpenTokenizer(OpenTokenizer),

crates/milli/src/vector/hf.rs

@@ -34,6 +34,30 @@ pub struct EmbedderOptions {
     pub model: String,
     pub revision: Option<String>,
     pub distribution: Option<DistributionShift>,
+    #[serde(default)]
+    pub pooling: OverridePooling,
+}
+
+#[derive(
+    Debug,
+    Clone,
+    Copy,
+    Default,
+    Hash,
+    PartialEq,
+    Eq,
+    serde::Deserialize,
+    serde::Serialize,
+    utoipa::ToSchema,
+    deserr::Deserr,
+)]
+#[deserr(rename_all = camelCase, deny_unknown_fields)]
+#[serde(rename_all = "camelCase")]
+pub enum OverridePooling {
+    UseModel,
+    ForceCls,
+    #[default]
+    ForceMean,
 }
 
 impl EmbedderOptions {
@@ -42,6 +66,7 @@ impl EmbedderOptions {
             model: "BAAI/bge-base-en-v1.5".to_string(),
             revision: Some("617ca489d9e86b49b8167676d8220688b99db36e".into()),
             distribution: None,
+            pooling: OverridePooling::UseModel,
         }
     }
 }
@@ -58,6 +83,7 @@ pub struct Embedder {
     tokenizer: Tokenizer,
     options: EmbedderOptions,
     dimensions: usize,
+    pooling: Pooling,
 }
 
 impl std::fmt::Debug for Embedder {
@@ -66,10 +92,62 @@ impl std::fmt::Debug for Embedder {
             .field("model", &self.options.model)
             .field("tokenizer", &self.tokenizer)
             .field("options", &self.options)
+            .field("pooling", &self.pooling)
             .finish()
     }
 }
 
+#[derive(Clone, Copy, serde::Deserialize)]
+struct PoolingConfig {
+    #[serde(default)]
+    pub pooling_mode_cls_token: bool,
+    #[serde(default)]
+    pub pooling_mode_mean_tokens: bool,
+    #[serde(default)]
+    pub pooling_mode_max_tokens: bool,
+    #[serde(default)]
+    pub pooling_mode_mean_sqrt_len_tokens: bool,
+    #[serde(default)]
+    pub pooling_mode_lasttoken: bool,
+}
+
+#[derive(Debug, Clone, Copy, Default)]
+pub enum Pooling {
+    #[default]
+    Mean,
+    Cls,
+    Max,
+    MeanSqrtLen,
+    LastToken,
+}
+impl Pooling {
+    fn override_with(&mut self, pooling: OverridePooling) {
+        match pooling {
+            OverridePooling::UseModel => {}
+            OverridePooling::ForceCls => *self = Pooling::Cls,
+            OverridePooling::ForceMean => *self = Pooling::Mean,
+        }
+    }
+}
+
+impl From<PoolingConfig> for Pooling {
+    fn from(value: PoolingConfig) -> Self {
+        if value.pooling_mode_cls_token {
+            Self::Cls
+        } else if value.pooling_mode_mean_tokens {
+            Self::Mean
+        } else if value.pooling_mode_lasttoken {
+            Self::LastToken
+        } else if value.pooling_mode_mean_sqrt_len_tokens {
+            Self::MeanSqrtLen
+        } else if value.pooling_mode_max_tokens {
+            Self::Max
+        } else {
+            Self::default()
+        }
+    }
+}
+
 impl Embedder {
     pub fn new(options: EmbedderOptions) -> std::result::Result<Self, NewEmbedderError> {
         let device = match candle_core::Device::cuda_if_available(0) {
@@ -83,7 +161,7 @@ impl Embedder {
             Some(revision) => Repo::with_revision(options.model.clone(), RepoType::Model, revision),
             None => Repo::model(options.model.clone()),
         };
-        let (config_filename, tokenizer_filename, weights_filename, weight_source) = {
+        let (config_filename, tokenizer_filename, weights_filename, weight_source, pooling) = {
             let api = Api::new().map_err(NewEmbedderError::new_api_fail)?;
             let api = api.repo(repo);
             let config = api.get("config.json").map_err(NewEmbedderError::api_get)?;
@@ -97,7 +175,38 @@ impl Embedder {
                     })
                     .map_err(NewEmbedderError::api_get)?
             };
-            (config, tokenizer, weights, source)
+            let pooling = match api.get("1_Pooling/config.json") {
+                Ok(pooling) => Some(pooling),
+                Err(hf_hub::api::sync::ApiError::RequestError(error))
+                    if matches!(*error, ureq::Error::Status(404, _,)) =>
+                {
+                    // ignore the error if the file simply doesn't exist
+                    None
+                }
+                Err(error) => return Err(NewEmbedderError::api_get(error)),
+            };
+            let mut pooling: Pooling = match pooling {
+                Some(pooling_filename) => {
+                    let pooling = std::fs::read_to_string(&pooling_filename).map_err(|inner| {
+                        NewEmbedderError::open_pooling_config(pooling_filename.clone(), inner)
+                    })?;
+
+                    let pooling: PoolingConfig =
+                        serde_json::from_str(&pooling).map_err(|inner| {
+                            NewEmbedderError::deserialize_pooling_config(
+                                options.model.clone(),
+                                pooling_filename,
+                                inner,
+                            )
+                        })?;
+                    pooling.into()
+                }
+                None => Pooling::default(),
+            };
+
+            pooling.override_with(options.pooling);
+
+            (config, tokenizer, weights, source, pooling)
         };
 
         let config = std::fs::read_to_string(&config_filename)
@@ -122,6 +231,8 @@ impl Embedder {
             },
         };
 
+        tracing::debug!(model = options.model, weight=?weight_source, pooling=?pooling, "model config");
+
         let model = BertModel::load(vb, &config).map_err(NewEmbedderError::load_model)?;
 
         if let Some(pp) = tokenizer.get_padding_mut() {
@@ -134,7 +245,7 @@ impl Embedder {
             tokenizer.with_padding(Some(pp));
         }
 
-        let mut this = Self { model, tokenizer, options, dimensions: 0 };
+        let mut this = Self { model, tokenizer, options, dimensions: 0, pooling };
 
         let embeddings = this
             .embed(vec!["test".into()])
@@ -168,17 +279,53 @@ impl Embedder {
             .forward(&token_ids, &token_type_ids, None)
             .map_err(EmbedError::model_forward)?;
 
-        // Apply some avg-pooling by taking the mean embedding value for all tokens (including padding)
-        let (_n_sentence, n_tokens, _hidden_size) =
-            embeddings.dims3().map_err(EmbedError::tensor_shape)?;
-
-        let embeddings = (embeddings.sum(1).map_err(EmbedError::tensor_value)? / (n_tokens as f64))
-            .map_err(EmbedError::tensor_shape)?;
+        let embeddings = Self::pooling(embeddings, self.pooling)?;
 
         let embeddings: Vec<Embedding> = embeddings.to_vec2().map_err(EmbedError::tensor_shape)?;
         Ok(embeddings)
     }
 
+    fn pooling(embeddings: Tensor, pooling: Pooling) -> Result<Tensor, EmbedError> {
+        match pooling {
+            Pooling::Mean => Self::mean_pooling(embeddings),
+            Pooling::Cls => Self::cls_pooling(embeddings),
+            Pooling::Max => Self::max_pooling(embeddings),
+            Pooling::MeanSqrtLen => Self::mean_sqrt_pooling(embeddings),
+            Pooling::LastToken => Self::last_token_pooling(embeddings),
+        }
+    }
+
+    fn cls_pooling(embeddings: Tensor) -> Result<Tensor, EmbedError> {
+        embeddings.get_on_dim(1, 0).map_err(EmbedError::tensor_value)
+    }
+
+    fn mean_sqrt_pooling(embeddings: Tensor) -> Result<Tensor, EmbedError> {
+        let (_n_sentence, n_tokens, _hidden_size) =
+            embeddings.dims3().map_err(EmbedError::tensor_shape)?;
+
+        (embeddings.sum(1).map_err(EmbedError::tensor_value)? / (n_tokens as f64).sqrt())
+            .map_err(EmbedError::tensor_shape)
+    }
+
+    fn mean_pooling(embeddings: Tensor) -> Result<Tensor, EmbedError> {
+        let (_n_sentence, n_tokens, _hidden_size) =
+            embeddings.dims3().map_err(EmbedError::tensor_shape)?;
+
+        (embeddings.sum(1).map_err(EmbedError::tensor_value)? / (n_tokens as f64))
+            .map_err(EmbedError::tensor_shape)
+    }
+
+    fn max_pooling(embeddings: Tensor) -> Result<Tensor, EmbedError> {
+        embeddings.max(1).map_err(EmbedError::tensor_shape)
+    }
+
+    fn last_token_pooling(embeddings: Tensor) -> Result<Tensor, EmbedError> {
+        let (_n_sentence, n_tokens, _hidden_size) =
+            embeddings.dims3().map_err(EmbedError::tensor_shape)?;
+
+        embeddings.get_on_dim(1, n_tokens - 1).map_err(EmbedError::tensor_value)
+    }
+
     pub fn embed_one(&self, text: &str) -> std::result::Result<Embedding, EmbedError> {
         let tokens = self.tokenizer.encode(text, true).map_err(EmbedError::tokenize)?;
         let token_ids = tokens.get_ids();
@@ -192,11 +339,8 @@ impl Embedder {
             .forward(&token_ids, &token_type_ids, None)
             .map_err(EmbedError::model_forward)?;
 
-        // Apply some avg-pooling by taking the mean embedding value for all tokens (including padding)
-        let (_n_sentence, n_tokens, _hidden_size) =
-            embeddings.dims3().map_err(EmbedError::tensor_shape)?;
-        let embedding = (embeddings.sum(1).map_err(EmbedError::tensor_value)? / (n_tokens as f64))
-            .map_err(EmbedError::tensor_shape)?;
+        let embedding = Self::pooling(embeddings, self.pooling)?;
+
         let embedding = embedding.squeeze(0).map_err(EmbedError::tensor_shape)?;
         let embedding: Embedding = embedding.to_vec1().map_err(EmbedError::tensor_shape)?;
         Ok(embedding)

crates/milli/src/vector/settings.rs

@@ -6,6 +6,7 @@ use roaring::RoaringBitmap;
 use serde::{Deserialize, Serialize};
 use utoipa::ToSchema;
 
+use super::hf::OverridePooling;
 use super::{ollama, openai, DistributionShift};
 use crate::prompt::{default_max_bytes, PromptData};
 use crate::update::Setting;
@@ -30,6 +31,10 @@ pub struct EmbeddingSettings {
     pub revision: Setting<String>,
     #[serde(default, skip_serializing_if = "Setting::is_not_set")]
     #[deserr(default)]
+    #[schema(value_type = Option<OverridePooling>)]
+    pub pooling: Setting<OverridePooling>,
+    #[serde(default, skip_serializing_if = "Setting::is_not_set")]
+    #[deserr(default)]
     #[schema(value_type = Option<String>)]
     pub api_key: Setting<String>,
     #[serde(default, skip_serializing_if = "Setting::is_not_set")]
@@ -164,6 +169,7 @@ impl SettingsDiff {
                     mut source,
                     mut model,
                     mut revision,
+                    mut pooling,
                     mut api_key,
                     mut dimensions,
                     mut document_template,
@@ -180,6 +186,7 @@ impl SettingsDiff {
                     source: new_source,
                     model: new_model,
                     revision: new_revision,
+                    pooling: new_pooling,
                     api_key: new_api_key,
                     dimensions: new_dimensions,
                     document_template: new_document_template,
@@ -210,6 +217,7 @@ impl SettingsDiff {
                         &source,
                         &mut model,
                         &mut revision,
+                        &mut pooling,
                         &mut dimensions,
                         &mut url,
                         &mut request,
@@ -225,6 +233,9 @@ impl SettingsDiff {
                 if revision.apply(new_revision) {
                     ReindexAction::push_action(&mut reindex_action, ReindexAction::FullReindex);
                 }
+                if pooling.apply(new_pooling) {
+                    ReindexAction::push_action(&mut reindex_action, ReindexAction::FullReindex);
+                }
                 if dimensions.apply(new_dimensions) {
                     match source {
                         // regenerate on dimensions change in OpenAI since truncation is supported
@@ -290,6 +301,7 @@ impl SettingsDiff {
                     source,
                     model,
                     revision,
+                    pooling,
                     api_key,
                     dimensions,
                     document_template,
@@ -338,6 +350,7 @@ fn apply_default_for_source(
     source: &Setting<EmbedderSource>,
     model: &mut Setting<String>,
     revision: &mut Setting<String>,
+    pooling: &mut Setting<OverridePooling>,
     dimensions: &mut Setting<usize>,
     url: &mut Setting<String>,
     request: &mut Setting<serde_json::Value>,
@@ -350,6 +363,7 @@ fn apply_default_for_source(
         Setting::Set(EmbedderSource::HuggingFace) => {
             *model = Setting::Reset;
             *revision = Setting::Reset;
+            *pooling = Setting::Reset;
             *dimensions = Setting::NotSet;
             *url = Setting::NotSet;
             *request = Setting::NotSet;
@@ -359,6 +373,7 @@ fn apply_default_for_source(
         Setting::Set(EmbedderSource::Ollama) => {
             *model = Setting::Reset;
             *revision = Setting::NotSet;
+            *pooling = Setting::NotSet;
             *dimensions = Setting::Reset;
             *url = Setting::NotSet;
             *request = Setting::NotSet;
@@ -368,6 +383,7 @@ fn apply_default_for_source(
         Setting::Set(EmbedderSource::OpenAi) | Setting::Reset => {
             *model = Setting::Reset;
             *revision = Setting::NotSet;
+            *pooling = Setting::NotSet;
             *dimensions = Setting::NotSet;
             *url = Setting::Reset;
             *request = Setting::NotSet;
@@ -377,6 +393,7 @@ fn apply_default_for_source(
         Setting::Set(EmbedderSource::Rest) => {
             *model = Setting::NotSet;
             *revision = Setting::NotSet;
+            *pooling = Setting::NotSet;
             *dimensions = Setting::Reset;
             *url = Setting::Reset;
             *request = Setting::Reset;
@@ -386,6 +403,7 @@ fn apply_default_for_source(
         Setting::Set(EmbedderSource::UserProvided) => {
             *model = Setting::NotSet;
             *revision = Setting::NotSet;
+            *pooling = Setting::NotSet;
             *dimensions = Setting::Reset;
             *url = Setting::NotSet;
             *request = Setting::NotSet;
@@ -419,6 +437,7 @@ impl EmbeddingSettings {
     pub const SOURCE: &'static str = "source";
     pub const MODEL: &'static str = "model";
     pub const REVISION: &'static str = "revision";
+    pub const POOLING: &'static str = "pooling";
     pub const API_KEY: &'static str = "apiKey";
     pub const DIMENSIONS: &'static str = "dimensions";
     pub const DOCUMENT_TEMPLATE: &'static str = "documentTemplate";
@@ -446,6 +465,7 @@ impl EmbeddingSettings {
                 &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama]
             }
             Self::REVISION => &[EmbedderSource::HuggingFace],
+            Self::POOLING => &[EmbedderSource::HuggingFace],
             Self::API_KEY => {
                 &[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest]
             }
@@ -500,6 +520,7 @@ impl EmbeddingSettings {
                 Self::SOURCE,
                 Self::MODEL,
                 Self::REVISION,
+                Self::POOLING,
                 Self::DOCUMENT_TEMPLATE,
                 Self::DOCUMENT_TEMPLATE_MAX_BYTES,
                 Self::DISTRIBUTION,
@@ -592,10 +613,12 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
                 model,
                 revision,
                 distribution,
+                pooling,
             }) => Self {
                 source: Setting::Set(EmbedderSource::HuggingFace),
                 model: Setting::Set(model),
                 revision: Setting::some_or_not_set(revision),
+                pooling: Setting::Set(pooling),
                 api_key: Setting::NotSet,
                 dimensions: Setting::NotSet,
                 document_template: Setting::Set(prompt.template),
@@ -617,6 +640,7 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
                 source: Setting::Set(EmbedderSource::OpenAi),
                 model: Setting::Set(embedding_model.name().to_owned()),
                 revision: Setting::NotSet,
+                pooling: Setting::NotSet,
                 api_key: Setting::some_or_not_set(api_key),
                 dimensions: Setting::some_or_not_set(dimensions),
                 document_template: Setting::Set(prompt.template),
@@ -638,6 +662,7 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
                 source: Setting::Set(EmbedderSource::Ollama),
                 model: Setting::Set(embedding_model),
                 revision: Setting::NotSet,
+                pooling: Setting::NotSet,
                 api_key: Setting::some_or_not_set(api_key),
                 dimensions: Setting::some_or_not_set(dimensions),
                 document_template: Setting::Set(prompt.template),
@@ -656,6 +681,7 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
                 source: Setting::Set(EmbedderSource::UserProvided),
                 model: Setting::NotSet,
                 revision: Setting::NotSet,
+                pooling: Setting::NotSet,
                 api_key: Setting::NotSet,
                 dimensions: Setting::Set(dimensions),
                 document_template: Setting::NotSet,
@@ -679,6 +705,7 @@ impl From<EmbeddingConfig> for EmbeddingSettings {
                 source: Setting::Set(EmbedderSource::Rest),
                 model: Setting::NotSet,
                 revision: Setting::NotSet,
+                pooling: Setting::NotSet,
                 api_key: Setting::some_or_not_set(api_key),
                 dimensions: Setting::some_or_not_set(dimensions),
                 document_template: Setting::Set(prompt.template),
@@ -701,6 +728,7 @@ impl From<EmbeddingSettings> for EmbeddingConfig {
             source,
             model,
             revision,
+            pooling,
             api_key,
             dimensions,
             document_template,
@@ -764,6 +792,9 @@ impl From<EmbeddingSettings> for EmbeddingConfig {
                     if let Some(revision) = revision.set() {
                         options.revision = Some(revision);
                     }
+                    if let Some(pooling) = pooling.set() {
+                        options.pooling = pooling;
+                    }
                     options.distribution = distribution.set();
                     this.embedder_options = super::EmbedderOptions::HuggingFace(options);
                 }