Add mean pooling strategy for Modernbert classifier (#616)

Author: kwnath

backends/candle/src/models/flash_modernbert.rs

@@ -260,7 +260,11 @@ impl FlashModernBertModel {
 
         let (pool, classifier) = match model_type {
             ModelType::Classifier => {
-                let pool = Pool::Cls;
+                let pool: Pool = config
+                    .classifier_pooling
+                    .as_deref()
+                    .and_then(|s| Pool::from_str(s).ok())
+                    .unwrap_or(Pool::Cls);
 
                 let classifier: Box<dyn ClassificationHead + Send> =
                     Box::new(ModernBertClassificationHead::load(vb.clone(), config)?);

backends/candle/src/models/modernbert.rs

@@ -7,6 +7,7 @@ use candle::{DType, Device, IndexOp, Module, Result, Tensor, D};
 use candle_nn::{Embedding, VarBuilder};
 use serde::Deserialize;
 use text_embeddings_backend_core::{Batch, ModelType, Pool};
+use std::str::FromStr;
 
 // https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/configuration_modernbert.py
 #[derive(Debug, Clone, PartialEq, Deserialize)]
@@ -484,7 +485,11 @@ impl ModernBertModel {
     pub fn load(vb: VarBuilder, config: &ModernBertConfig, model_type: ModelType) -> Result<Self> {
         let (pool, classifier) = match model_type {
             ModelType::Classifier => {
-                let pool = Pool::Cls;
+                let pool: Pool = config
+                    .classifier_pooling
+                    .as_deref()
+                    .and_then(|s| Pool::from_str(s).ok())
+                    .unwrap_or(Pool::Cls);
 
                 let classifier: Box<dyn ClassificationHead + Send> =
                     Box::new(ModernBertClassificationHead::load(vb.clone(), config)?);

backends/candle/tests/snapshots/test_modernbert__modernbert_classification_mean_pooling.snap.new

@@ -0,0 +1,7 @@
+---
+source: backends/candle/tests/test_modernbert.rs
+assertion_line: 229
+expression: predictions_single
+---
+- - -0.30617672
+

backends/candle/tests/test_modernbert.rs

@@ -202,3 +202,35 @@ fn test_modernbert_classification() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[serial_test::serial]
+fn test_modernbert_classification_mean_pooling() -> Result<()> {
+    let model_root = download_artifacts("tomaarsen/reranker-ModernBERT-large-gooaq-bce", None)?;
+    let tokenizer = load_tokenizer(&model_root)?;
+    let backend = CandleBackend::new(&model_root, "float32".to_string(), ModelType::Classifier)?;
+
+    let input_single = batch(
+        vec![tokenizer
+            .encode(("What is Deep Learning?", "Deep Learning is not..."), true)
+            .unwrap()],
+        [0].to_vec(),
+        vec![],
+    );
+
+    let predictions: Vec<Vec<f32>> = backend
+        .predict(input_single)?
+        .into_iter()
+        .map(|(_, v)| v)
+        .collect();
+    let predictions_single = SnapshotScores::from(predictions);
+
+    let matcher = relative_matcher();
+    insta::assert_yaml_snapshot!(
+        "modernbert_classification_mean_pooling",
+        predictions_single,
+        &matcher
+    );
+
+    Ok(())
+}

backends/core/src/lib.rs

@@ -78,6 +78,23 @@ impl fmt::Display for Pool {
     }
 }
 
+impl std::str::FromStr for Pool {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.trim().to_lowercase().as_str() {
+            "cls" => Ok(Pool::Cls),
+            "mean" => Ok(Pool::Mean),
+            "splade" => Ok(Pool::Splade),
+            "last_token" => Ok(Pool::LastToken),
+            _ => Err(format!(
+                "Invalid pooling method '{}'. Valid options: cls, mean, splade, last_token", 
+                s
+            )),
+        }
+    }
+}
+
 #[derive(Debug, Error, Clone)]
 pub enum BackendError {
     #[error("No backend found")]