Support modular Sentence Transformers cross-encoder rerankers (e.g. ettin-reranker)

backends/candle/src/lib.rs

@@ -7,7 +7,7 @@ mod layers;
 mod models;
 
 use anyhow::Context;
-use candle::{DType, Device};
+use candle::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
 use nohash_hasher::BuildNoHashHasher;
 use serde::{de::Deserializer, Deserialize};
@@ -26,7 +26,8 @@ use crate::models::{
     DistilBertConfig, DistilBertModel, GTEConfig, GTEModel, Gemma3Config, Gemma3Model,
     JinaBertModel, JinaCodeBertModel, LlamaConfig, MPNetConfig, MPNetModel, MistralConfig, Model,
     ModernBertConfig, ModernBertModel, NomicBertModel, NomicConfig, Pplx1Config, Pplx1Model,
-    Qwen2Config, Qwen3Config, Qwen3Model,
+    PredictionHeadLayerNorm, PredictionHeadModule, PredictionHeadModuleConfig, Qwen2Config,
+    Qwen3Config, Qwen3Model,
 };
 #[cfg(feature = "cuda")]
 use crate::models::{
@@ -157,6 +158,7 @@ pub struct CandleBackend {
     device: Device,
     model: Box<dyn Model + Send>,
     dense_layers: Vec<Box<dyn DenseLayer + Send>>,
+    prediction_head: Option<Vec<Box<dyn PredictionHeadModule + Send>>>,
 }
 
 impl CandleBackend {
@@ -165,6 +167,25 @@ impl CandleBackend {
         dtype: String,
         model_type: ModelType,
         dense_paths: Option<Vec<String>>,
+    ) -> Result<Self, BackendError> {
+        Self::load(model_path, dtype, model_type, dense_paths, false)
+    }
+
+    pub fn new_with_post_pooling_prediction(
+        model_path: &Path,
+        dtype: String,
+        model_type: ModelType,
+        dense_paths: Option<Vec<String>>,
+    ) -> Result<Self, BackendError> {
+        Self::load(model_path, dtype, model_type, dense_paths, true)
+    }
+
+    fn load(
+        model_path: &Path,
+        dtype: String,
+        model_type: ModelType,
+        dense_paths: Option<Vec<String>>,
+        use_post_pooling_prediction: bool,
     ) -> Result<Self, BackendError> {
         // Default files
         let default_safetensors = model_path.join("model.safetensors");
@@ -577,7 +598,74 @@ impl CandleBackend {
         };
 
         let mut dense_layers = Vec::new();
-        if let Some(dense_paths) = dense_paths {
+        let mut prediction_head = None;
+        if use_post_pooling_prediction {
+            // Modular Sentence Transformers reranker head: the post-pooling
+            // `Dense`/`LayerNorm` modules that produce the rerank score.
+            let module_paths = dense_paths
+                .filter(|paths| !paths.is_empty())
+                .ok_or_else(|| {
+                    BackendError::Start(
+                        "Post-pooling prediction requires at least one module".to_string(),
+                    )
+                })?;
+
+            tracing::info!(
+                "Loading post-pooling prediction module/s from path/s: {module_paths:?}"
+            );
+
+            let mut modules = Vec::new();
+            for module_path in module_paths.iter() {
+                let module_safetensors =
+                    model_path.join(format!("{module_path}/model.safetensors"));
+                let module_pytorch = model_path.join(format!("{module_path}/pytorch_model.bin"));
+
+                if !(module_safetensors.exists() || module_pytorch.exists()) {
+                    return Err(BackendError::Start(format!(
+                        "Post-pooling prediction module files not found for path: {module_path}"
+                    )));
+                }
+
+                let module_config_path = model_path.join(format!("{module_path}/config.json"));
+                let module_config_str =
+                    std::fs::read_to_string(&module_config_path).map_err(|err| {
+                        BackendError::Start(format!(
+                            "Unable to read `{module_path}/config.json` file: {err:?}",
+                        ))
+                    })?;
+                let module_config: PredictionHeadModuleConfig =
+                    serde_json::from_str(&module_config_str).map_err(|err| {
+                        BackendError::Start(format!(
+                            "Unable to parse `{module_path}/config.json`: {err:?}",
+                        ))
+                    })?;
+
+                let module_vb = if module_safetensors.exists() {
+                    unsafe {
+                        VarBuilder::from_mmaped_safetensors(&[module_safetensors], dtype, &device)
+                    }
+                    .s()?
+                } else {
+                    VarBuilder::from_pth(&module_pytorch, dtype, &device).s()?
+                };
+
+                let module = match module_config {
+                    PredictionHeadModuleConfig::Dense(config) => {
+                        Box::new(Dense::load(module_vb, &config).s()?)
+                            as Box<dyn PredictionHeadModule + Send>
+                    }
+                    PredictionHeadModuleConfig::LayerNorm(config) => {
+                        Box::new(PredictionHeadLayerNorm::load(module_vb, &config).s()?)
+                            as Box<dyn PredictionHeadModule + Send>
+                    }
+                };
+                modules.push(module);
+
+                tracing::info!("Loaded post-pooling prediction module from path: {module_path}");
+            }
+
+            prediction_head = Some(modules);
+        } else if let Some(dense_paths) = dense_paths {
             if !dense_paths.is_empty() {
                 tracing::info!("Loading Dense module/s from path/s: {dense_paths:?}");
 
@@ -632,8 +720,36 @@ impl CandleBackend {
             device,
             model: model?,
             dense_layers,
+            prediction_head,
         })
     }
+
+    fn predict_from_pooled_embeddings(
+        &self,
+        batch: Batch,
+        prediction_head: &[Box<dyn PredictionHeadModule + Send>],
+    ) -> Result<Tensor, BackendError> {
+        let batch_size = batch.len();
+        let (Some(mut pooled_embeddings), _) = self.model.embed(batch).e()? else {
+            return Err(BackendError::Inference(
+                "prediction head did not receive pooled embeddings".to_string(),
+            ));
+        };
+        for module in prediction_head {
+            pooled_embeddings = module.forward(&pooled_embeddings).e()?;
+        }
+
+        match pooled_embeddings.dims() {
+            [rows, cols] if *rows == batch_size && *cols == 1 => {}
+            shape => {
+                return Err(BackendError::Inference(format!(
+                    "prediction head returned shape {shape:?}; expected [{batch_size}, 1]"
+                )));
+            }
+        }
+
+        Ok(pooled_embeddings)
+    }
 }
 
 impl Backend for CandleBackend {
@@ -711,7 +827,11 @@ impl Backend for CandleBackend {
     fn predict(&self, batch: Batch) -> Result<Predictions, BackendError> {
         let batch_size = batch.len();
 
-        let results = self.model.predict(batch).e()?;
+        let results = if let Some(prediction_head) = &self.prediction_head {
+            self.predict_from_pooled_embeddings(batch, prediction_head)?
+        } else {
+            self.model.predict(batch).e()?
+        };
 
         let results = results.to_dtype(DType::F32).e()?.to_vec2().e()?;
 

backends/candle/src/models/dense.rs

@@ -1,4 +1,4 @@
-use crate::layers::Linear;
+use crate::layers::{LayerNorm, Linear};
 use candle::{Result, Tensor};
 use candle_nn::VarBuilder;
 use serde::Deserialize;
@@ -12,17 +12,27 @@ pub enum DenseActivation {
     #[serde(rename = "torch.nn.modules.linear.Identity")]
     /// e.g. https://huggingface.co/NovaSearch/stella_en_400M_v5/blob/main/2_Dense/config.json
     Identity,
+    #[serde(rename = "torch.nn.modules.activation.GELU")]
+    Gelu,
 }
 
 impl DenseActivation {
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
         match self {
             Self::Tanh => x.tanh(),
             Self::Identity => Ok(x.clone()),
+            Self::Gelu => x.gelu(),
         }
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+#[serde(untagged)]
+pub enum PredictionHeadModuleConfig {
+    Dense(DenseConfig),
+    LayerNorm(LayerNormConfig),
+}
+
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct DenseConfig {
     in_features: usize,
@@ -31,10 +41,20 @@ pub struct DenseConfig {
     activation_function: Option<DenseActivation>,
 }
 
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct LayerNormConfig {
+    dimension: usize,
+}
+
 pub trait DenseLayer {
     fn forward(&self, hidden_states: &Tensor) -> Result<Tensor>;
 }
 
+/// A post-pooling scoring-head module (`Dense`/`LayerNorm`) of a modular Sentence Transformers reranker.
+pub trait PredictionHeadModule {
+    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor>;
+}
+
 #[derive(Debug)]
 pub struct Dense {
     linear: Linear,
@@ -73,3 +93,29 @@ impl DenseLayer for Dense {
         self.activation.forward(&hidden_states)
     }
 }
+
+impl PredictionHeadModule for Dense {
+    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
+        DenseLayer::forward(self, hidden_states)
+    }
+}
+
+#[derive(Debug)]
+pub struct PredictionHeadLayerNorm {
+    layer_norm: LayerNorm,
+}
+
+impl PredictionHeadLayerNorm {
+    pub fn load(vb: VarBuilder, config: &LayerNormConfig) -> Result<Self> {
+        Ok(Self {
+            layer_norm: LayerNorm::load(vb.clone(), config.dimension, 1e-5)
+                .or_else(|_| LayerNorm::load(vb.pp("norm"), config.dimension, 1e-5))?,
+        })
+    }
+}
+
+impl PredictionHeadModule for PredictionHeadLayerNorm {
+    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
+        self.layer_norm.forward(hidden_states, None)
+    }
+}

backends/candle/src/models/flash_modernbert.rs

@@ -298,9 +298,9 @@ impl FlashModernBertModel {
 
         for use_local_attention in [true, false] {
             let rope_theta = if use_local_attention {
-                config.local_rope_theta
+                config.local_rope_theta()?
             } else {
-                config.global_rope_theta
+                config.global_rope_theta()?
             };
 
             let max_position_embeddings = config.max_position_embeddings;

backends/candle/src/models/mod.rs

@@ -56,7 +56,10 @@ mod flash_qwen3;
 
 pub use bert::{BertConfig, BertModel, PositionEmbeddingType};
 pub use debertav2::{DebertaV2Config, DebertaV2Model};
-pub use dense::{Dense, DenseConfig, DenseLayer};
+pub use dense::{
+    Dense, DenseConfig, DenseLayer, PredictionHeadLayerNorm, PredictionHeadModule,
+    PredictionHeadModuleConfig,
+};
 pub use distilbert::{DistilBertConfig, DistilBertModel};
 pub use gemma3::{Gemma3Config, Gemma3Model};
 pub use gte::{GTEConfig, GTEModel};

backends/candle/src/models/modernbert.rs

@@ -23,16 +23,17 @@ pub struct ModernBertConfig {
     pub norm_eps: f64,
     pub norm_bias: bool,
     pub pad_token_id: usize,
-    pub eos_token_id: usize,
-    pub bos_token_id: usize,
+    pub eos_token_id: Option<usize>,
+    pub bos_token_id: Option<usize>,
     pub cls_token_id: usize,
     pub sep_token_id: usize,
-    pub global_rope_theta: f64,
+    pub global_rope_theta: Option<f64>,
     pub attention_bias: bool,
     pub attention_dropout: f64,
     pub global_attn_every_n_layers: usize,
     pub local_attention: usize,
-    pub local_rope_theta: f64,
+    pub local_rope_theta: Option<f64>,
+    pub rope_parameters: Option<ModernBertRopeParameters>,
     pub embedding_dropout: Option<f64>,
     pub mlp_bias: Option<bool>,
     pub mlp_dropout: Option<f64>,
@@ -48,6 +49,49 @@ pub struct ModernBertConfig {
     pub num_labels: Option<usize>,
 }
 
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct ModernBertRopeParameters {
+    full_attention: Option<ModernBertRopeParametersEntry>,
+    sliding_attention: Option<ModernBertRopeParametersEntry>,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct ModernBertRopeParametersEntry {
+    rope_theta: Option<f64>,
+}
+
+impl ModernBertConfig {
+    pub fn global_rope_theta(&self) -> Result<f64> {
+        if let Some(theta) = self.global_rope_theta {
+            return Ok(theta);
+        }
+        if let Some(theta) = self
+            .rope_parameters
+            .as_ref()
+            .and_then(|params| params.full_attention.as_ref())
+            .and_then(|params| params.rope_theta)
+        {
+            return Ok(theta);
+        }
+        candle::bail!("`global_rope_theta` or `rope_parameters.full_attention.rope_theta` must be defined in `config.json`")
+    }
+
+    pub fn local_rope_theta(&self) -> Result<f64> {
+        if let Some(theta) = self.local_rope_theta {
+            return Ok(theta);
+        }
+        if let Some(theta) = self
+            .rope_parameters
+            .as_ref()
+            .and_then(|params| params.sliding_attention.as_ref())
+            .and_then(|params| params.rope_theta)
+        {
+            return Ok(theta);
+        }
+        candle::bail!("`local_rope_theta` or `rope_parameters.sliding_attention.rope_theta` must be defined in `config.json`")
+    }
+}
+
 #[derive(Debug)]
 pub struct ModernBertEmbeddings {
     tok_embeddings: Embedding,
@@ -521,13 +565,13 @@ impl ModernBertModel {
 
         let global_inv_freqs = get_inv_freqs(
             attention_head_size,
-            config.global_rope_theta as f32,
+            config.global_rope_theta()? as f32,
             vb.device(),
             None,
         )?;
         let local_inv_freqs = get_inv_freqs(
             attention_head_size,
-            config.local_rope_theta as f32,
+            config.local_rope_theta()? as f32,
             vb.device(),
             None,
         )?;