Add Dense, DenseLayer and DenseConfig to handle 2_Dense/

Required for some models as e.g.
https://huggingface.co/sentence-transformers/LaBSE

Author: alvarobartt

backends/candle/src/lib.rs

@@ -11,9 +11,10 @@ use crate::compute_cap::{
     compatible_compute_cap, get_compile_compute_cap, get_runtime_compute_cap,
 };
 use crate::models::{
-    BertConfig, BertModel, DistilBertConfig, DistilBertModel, GTEConfig, GTEModel, JinaBertModel,
-    JinaCodeBertModel, MPNetConfig, MPNetModel, MistralConfig, Model, ModernBertConfig,
-    ModernBertModel, NomicBertModel, NomicConfig, Qwen2Config, Qwen3Config, Qwen3Model,
+    BertConfig, BertModel, Dense, DenseConfig, DenseLayer, DistilBertConfig, DistilBertModel,
+    GTEConfig, GTEModel, JinaBertModel, JinaCodeBertModel, MPNetConfig, MPNetModel, MistralConfig,
+    Model, ModernBertConfig, ModernBertModel, NomicBertModel, NomicConfig, Qwen2Config,
+    Qwen3Config, Qwen3Model,
 };
 #[cfg(feature = "cuda")]
 use crate::models::{
@@ -114,6 +115,7 @@ enum Config {
 pub struct CandleBackend {
     device: Device,
     model: Box<dyn Model + Send>,
+    dense: Option<Box<dyn DenseLayer + Send>>,
 }
 
 impl CandleBackend {
@@ -468,9 +470,35 @@ impl CandleBackend {
             }
         };
 
+        // If `2_Dense/model.safetensors` is amongst the downloaded artifacts, then create a Linear
+        // layer from the VarBuilder using `candle` to provide it as an extra `Dense` layer to the
+        // `CandleBackend`, otherwise leave it as None
+        let dense = if model_path.join("2_Dense/model.safetensors").exists() {
+            let dense_config_path = model_path.join("2_Dense/config.json");
+
+            // Load dense config
+            let dense_config_str = std::fs::read_to_string(&dense_config_path).map_err(|err| {
+                BackendError::Start(format!("Unable to read dense config file: {err:?}"))
+            })?;
+            let dense_config: DenseConfig =
+                serde_json::from_str(&dense_config_str).map_err(|err| {
+                    BackendError::Start(format!("Unable to parse dense config: {err:?}"))
+                })?;
+
+            let dense_path = model_path.join("2_Dense/model.safetensors");
+            let dense_vb =
+                unsafe { VarBuilder::from_mmaped_safetensors(&[dense_path], dtype, &device) }
+                    .s()?;
+
+            Some(Box::new(Dense::load(dense_vb, &dense_config).s()?) as Box<dyn DenseLayer + Send>)
+        } else {
+            None
+        };
+
         Ok(Self {
             device,
             model: model?,
+            dense: dense,
         })
     }
 }
@@ -507,6 +535,19 @@ impl Backend for CandleBackend {
         // Run forward
         let (pooled_embeddings, raw_embeddings) = self.model.embed(batch).e()?;
 
+        // Apply dense layer if available
+        let pooled_embeddings = match pooled_embeddings {
+            None => None,
+            Some(pooled_embeddings) => {
+                let pooled_embeddings = if let Some(ref dense) = self.dense {
+                    dense.forward(&pooled_embeddings).e()?
+                } else {
+                    pooled_embeddings
+                };
+                Some(pooled_embeddings)
+            }
+        };
+
         // Device => Host data transfer
         let pooled_embeddings = match pooled_embeddings {
             None => vec![],
@@ -540,6 +581,7 @@ impl Backend for CandleBackend {
         let batch_size = batch.len();
 
         let results = self.model.predict(batch).e()?;
+
         let results = results.to_dtype(DType::F32).e()?.to_vec2().e()?;
 
         let mut predictions =

backends/candle/src/models/dense.rs

@@ -0,0 +1,48 @@
+use crate::layers::Linear;
+use candle::{Result, Tensor};
+use candle_nn::VarBuilder;
+use serde::Deserialize;
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct DenseConfig {
+    in_features: usize,
+    out_features: usize,
+    bias: bool,
+    #[allow(unused)]
+    activation_function: Option<String>,
+}
+
+pub trait DenseLayer {
+    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor>;
+}
+
+#[derive(Debug)]
+pub struct Dense {
+    linear: Linear,
+    span: tracing::Span,
+}
+
+impl Dense {
+    pub fn load(vb: VarBuilder, config: &DenseConfig) -> Result<Self> {
+        let dense_weight = vb.get((config.out_features, config.in_features), "linear.weight")?;
+        let dense_bias = if config.bias {
+            Some(vb.get(config.out_features, "linear.bias")?)
+        } else {
+            None
+        };
+
+        let linear = Linear::new(dense_weight, dense_bias, None);
+
+        Ok(Self {
+            linear,
+            span: tracing::span!(tracing::Level::TRACE, "dense"),
+        })
+    }
+}
+
+impl DenseLayer for Dense {
+    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
+        let _enter = self.span.enter();
+        self.linear.forward(hidden_states)?.tanh()
+    }
+}

backends/candle/src/models/mod.rs

@@ -5,6 +5,7 @@ extern crate intel_mkl_src;
 extern crate accelerate_src;
 
 mod bert;
+mod dense;
 mod distilbert;
 mod jina;
 mod jina_code;
@@ -49,6 +50,7 @@ mod qwen3;
 
 pub use bert::{BertConfig, BertModel, PositionEmbeddingType};
 use candle::{Result, Tensor};
+pub use dense::{Dense, DenseConfig, DenseLayer};
 pub use distilbert::{DistilBertConfig, DistilBertModel};
 #[allow(unused_imports)]
 pub use gte::{GTEClassificationHead, GTEConfig, GTEModel, GTEMLP};

core/src/download.rs

@@ -37,6 +37,14 @@ pub async fn download_artifacts(api: &ApiRepo, pool_config: bool) -> Result<Path
         err
     });
 
+    // Try to download the `2_Dense/config.json`
+    if let Ok(_) = download_dense_config(api).await {
+        // If `2_Dense/config.json` is there, then try to download the `model.safetensors`
+        if let Err(err) = download_dense_safetensors(api).await {
+            tracing::warn!("Failed to download dense safetensors: {err}");
+        }
+    }
+
     tracing::info!("Downloading `config.json`");
     api.get("config.json").await?;
 
@@ -55,6 +63,20 @@ pub async fn download_pool_config(api: &ApiRepo) -> Result<PathBuf, ApiError> {
     Ok(pool_config_path)
 }
 
+#[instrument(skip_all)]
+pub async fn download_dense_config(api: &ApiRepo) -> Result<PathBuf, ApiError> {
+    tracing::info!("Downloading `2_Dense/config.json`");
+    let dense_config_path = api.get("2_Dense/config.json").await?;
+    Ok(dense_config_path)
+}
+
+#[instrument(skip_all)]
+pub async fn download_dense_safetensors(api: &ApiRepo) -> Result<PathBuf, ApiError> {
+    tracing::info!("Downloading `2_Dense/model.safetensors`");
+    let dense_safetensors_path = api.get("2_Dense/model.safetensors").await?;
+    Ok(dense_safetensors_path)
+}
+
 #[instrument(skip_all)]
 pub async fn download_st_config(api: &ApiRepo) -> Result<PathBuf, ApiError> {
     // Try default path

router/src/lib.rs

@@ -405,6 +405,7 @@ fn get_backend_model_type(
             }
         }
     };
+
     Ok(text_embeddings_backend::ModelType::Embedding(pool))
 }