Adding missing head. prefix in the weight name in ModernBertClassificationHead (#591)

Author: kozistr

README.md

@@ -101,6 +101,7 @@ Below are some examples of the currently supported models:
 | Re-Ranking         | XLM-RoBERTa | [BAAI/bge-reranker-large](https://huggingface.co/BAAI/bge-reranker-large)                                       |
 | Re-Ranking         | XLM-RoBERTa | [BAAI/bge-reranker-base](https://huggingface.co/BAAI/bge-reranker-base)                                         |
 | Re-Ranking         | GTE         | [Alibaba-NLP/gte-multilingual-reranker-base](https://huggingface.co/Alibaba-NLP/gte-multilingual-reranker-base) |
+| Re-Ranking         | ModernBert  | [Alibaba-NLP/gte-reranker-modernbert-base](https://huggingface.co/Alibaba-NLP/gte-reranker-modernbert-base) |
 | Sentiment Analysis | RoBERTa     | [SamLowe/roberta-base-go_emotions](https://huggingface.co/SamLowe/roberta-base-go_emotions)                     |
 
 ### Docker

backends/candle/src/models/modernbert.rs

@@ -412,27 +412,28 @@ pub struct ModernBertClassificationHead {
 impl ModernBertClassificationHead {
     pub(crate) fn load(vb: VarBuilder, config: &ModernBertConfig) -> Result<Self> {
         let dense_weight = vb
-            .pp("dense")
+            .pp("head.dense")
             .get((config.hidden_size, config.hidden_size), "weight")?;
-        let dense_bias = vb.pp("dense").get(config.hidden_size, "bias").ok();
         let dense = Linear::new(
             dense_weight,
-            dense_bias,
+            None,
             Some(config.classifier_activation.clone()),
         );
 
-        let norm =
-            LayerNormNoBias::load(vb.pp("norm"), config.hidden_size, config.norm_eps as f32)?;
+        let norm = LayerNormNoBias::load(
+            vb.pp("head.norm"),
+            config.hidden_size,
+            config.norm_eps as f32,
+        )?;
 
-        let classifier_weight = vb.pp("dense").get(
+        let classifier_weight = vb.pp("classifier").get(
             (config.num_labels.unwrap_or(1), config.hidden_size),
             "weight",
         )?;
         let classifier_bias = vb
-            .pp("dense")
-            .get(config.num_labels.unwrap_or(1), "bias")
-            .ok();
-        let classifier = Linear::new(classifier_weight, classifier_bias, None);
+            .pp("classifier")
+            .get(config.num_labels.unwrap_or(1), "bias")?;
+        let classifier = Linear::new(classifier_weight, Some(classifier_bias), None);
 
         Ok(Self {
             dense,

backends/candle/tests/snapshots/test_modernbert__modernbert_classification_single.snap

@@ -0,0 +1,5 @@
+---
+source: backends/candle/tests/test_modernbert.rs
+expression: predictions_single
+---
+- - 2.2585099

backends/candle/tests/test_modernbert.rs

@@ -2,7 +2,9 @@ mod common;
 
 use crate::common::{sort_embeddings, SnapshotEmbeddings};
 use anyhow::Result;
-use common::{batch, cosine_matcher, download_artifacts, load_tokenizer};
+use common::{
+    batch, cosine_matcher, download_artifacts, load_tokenizer, relative_matcher, SnapshotScores,
+};
 use text_embeddings_backend_candle::CandleBackend;
 use text_embeddings_backend_core::{Backend, ModelType, Pool};
 
@@ -135,3 +137,36 @@ fn test_mini_pooled_raw() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[serial_test::serial]
+fn test_modernbert_classification() -> Result<()> {
+    let model_root = download_artifacts("Alibaba-NLP/gte-reranker-modernbert-base", None).unwrap();
+    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_single",
+        predictions_single,
+        &matcher
+    );
+
+    Ok(())
+}