feat: support pplx-embed-v1 (#824)

Co-authored-by: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>

Author: mkrimmel-pplx

backends/candle/src/lib.rs

@@ -25,8 +25,8 @@ use crate::models::{
     BertConfig, BertModel, DebertaV2Config, DebertaV2Model, Dense, DenseConfig, DenseLayer,
     DistilBertConfig, DistilBertModel, GTEConfig, GTEModel, Gemma3Config, Gemma3Model,
     JinaBertModel, JinaCodeBertModel, LlamaConfig, MPNetConfig, MPNetModel, MistralConfig, Model,
-    ModernBertConfig, ModernBertModel, NomicBertModel, NomicConfig, Qwen2Config, Qwen3Config,
-    Qwen3Model,
+    ModernBertConfig, ModernBertModel, NomicBertModel, NomicConfig, Pplx1Config, Pplx1Model,
+    Qwen2Config, Qwen3Config, Qwen3Model,
 };
 #[cfg(feature = "cuda")]
 use crate::models::{
@@ -135,16 +135,17 @@ enum Config {
     Gte(GTEConfig),
     #[serde(rename = "mpnet")]
     MPNet(MPNetConfig),
-    #[allow(dead_code)]
+    #[allow(dead_code)] // NOTE: As it's only used when `cuda` feature is enabled
     Mistral(MistralConfig),
     #[serde(rename(deserialize = "modernbert"))]
     ModernBert(ModernBertConfig),
     #[serde(rename(deserialize = "nomic_bert"))]
     NomicBert(NomicConfig),
-    #[allow(dead_code)]
+    #[allow(dead_code)] // NOTE: As it's only used when `cuda` feature is enabled
     Qwen2(Qwen2Config),
-    #[allow(dead_code)]
     Qwen3(Qwen3Config),
+    #[serde(rename(deserialize = "bidirectional_pplx_qwen3"))]
+    Pplx1(Pplx1Config),
     Roberta(BertConfig),
     XlmRoberta(BertConfig),
     #[allow(dead_code)]
@@ -319,6 +320,7 @@ impl CandleBackend {
                 ))
             }
             (Config::Gemma3(config), Device::Cpu | Device::Metal(_)) => {
+                // TODO(alvarobartt): Enable Flash Attention with BF16 once supported on Metal
                 if dtype != DType::F32 {
                     Err(BackendError::Start(
                         "Gemma3 is only supported in fp32 precision".to_string(),
@@ -362,6 +364,17 @@ impl CandleBackend {
                 tracing::info!("Starting Qwen3 model on {:?}", device);
                 Ok(Box::new(Qwen3Model::load(vb, &config, model_type).s()?))
             }
+            (Config::Pplx1(config), Device::Cpu | Device::Metal(_)) => {
+                // TODO(alvarobartt): Enable Flash Attention with BF16 once supported on Metal
+                if dtype != DType::F32 {
+                    Err(BackendError::Start(
+                        "Pplx1 is only supported in fp32 precision".to_string(),
+                    ))
+                } else {
+                    tracing::info!("Starting Pplx1 model on {:?}", device);
+                    Ok(Box::new(Pplx1Model::load(vb, &config, model_type).s()?))
+                }
+            }
             #[cfg(feature = "cuda")]
             (Config::Bert(config), Device::Cuda(_)) => {
                 if dtype == DType::F16 && use_flash_attn(&[FlashAttn::V1, FlashAttn::V2]) {
@@ -440,6 +453,7 @@ impl CandleBackend {
             }
             #[cfg(feature = "cuda")]
             (Config::Gemma3(config), Device::Cuda(_)) => {
+                // TODO(alvarobartt): Enable Flash Attention with BF16 once supported on CUDA
                 if dtype != DType::F32 {
                     Err(BackendError::Start(
                         "Gemma3 is only supported in fp32 precision".to_string(),
@@ -520,6 +534,18 @@ impl CandleBackend {
                 }
             }
             #[cfg(feature = "cuda")]
+            (Config::Pplx1(config), Device::Cuda(_)) => {
+                // TODO(alvarobartt): Enable Flash Attention with BF16 once supported on CUDA
+                if dtype != DType::F32 {
+                    Err(BackendError::Start(
+                        "Pplx1 is only supported in fp32 precision".to_string(),
+                    ))
+                } else {
+                    tracing::info!("Starting Pplx1 model on {:?}", device);
+                    Ok(Box::new(Pplx1Model::load(vb, &config, model_type).s()?))
+                }
+            }
+            #[cfg(feature = "cuda")]
             (Config::Llama(config), Device::Cuda(_)) => match config.rope_scaling {
                 Some(_) => Err(BackendError::Start(
                     "Rope scaling is not supported for FlashLlama yet".to_string(),

backends/candle/src/models/flash_mistral.rs

@@ -24,8 +24,6 @@ struct MistralAttention {
 
 impl MistralAttention {
     pub fn load(vb: VarBuilder, config: &MistralConfig) -> Result<Self> {
-        let window_size_left = config.sliding_window;
-        let use_bidirectional_attention = config.use_bidirectional_attention.unwrap_or(false);
         let num_attention_heads = config.num_attention_heads;
         let attention_head_size = config.hidden_size / config.num_attention_heads;
         let num_key_value_heads = config.num_key_value_heads;
@@ -55,8 +53,8 @@ impl MistralAttention {
         Ok(Self {
             qkv_linear,
             o_proj,
-            window_size_left,
-            use_bidirectional_attention,
+            window_size_left: config.sliding_window,
+            use_bidirectional_attention: config.use_bidirectional_attention.unwrap_or(false),
             num_attention_heads,
             num_key_value_heads,
             attention_head_size,

backends/candle/src/models/flash_qwen3.rs

@@ -162,6 +162,7 @@ impl Qwen3Attention {
             max_s,
             max_s,
             self.softmax_scale,
+            // NOTE: When `use_bidirectional_attention=true` then that implies that `causal=false`
             !self.use_bidirectional_attention,
             None,
             None,

backends/candle/src/models/mod.rs

@@ -20,6 +20,7 @@ mod mistral;
 mod modernbert;
 mod mpnet;
 mod nomic;
+mod pplx1;
 mod qwen2;
 mod qwen3;
 
@@ -66,6 +67,7 @@ pub use mistral::MistralConfig;
 pub use modernbert::{ModernBertConfig, ModernBertModel};
 pub use mpnet::{MPNetConfig, MPNetModel};
 pub use nomic::{NomicBertModel, NomicConfig};
+pub use pplx1::{Pplx1Config, Pplx1Model};
 pub use qwen2::Qwen2Config;
 pub use qwen3::{Qwen3Config, Qwen3Model};
 

backends/candle/src/models/pplx1.rs

@@ -0,0 +1,59 @@
+use crate::models::{Model, Qwen3Config, Qwen3Model};
+use candle::{Result, Tensor};
+use candle_nn::VarBuilder;
+use text_embeddings_backend_core::{Batch, ModelType, Pool};
+
+pub type Pplx1Config = Qwen3Config;
+
+pub struct Pplx1Model {
+    inner: Qwen3Model,
+}
+
+impl Pplx1Model {
+    pub fn load(vb: VarBuilder, config: &Pplx1Config, model_type: ModelType) -> Result<Self> {
+        match model_type {
+            ModelType::Classifier => {
+                candle::bail!("`classifier` model type is not supported for Pplx1")
+            }
+            ModelType::Embedding(ref pool) => {
+                if pool != &Pool::Mean {
+                    candle::bail!("Pplx1 only supports mean pooling, got {:?}", pool);
+                }
+            }
+        };
+
+        // NOTE: Qwen3 but the `config` contains `use_bidirectional_attention=true`
+        let inner = Qwen3Model::load(vb, config, model_type)?;
+
+        Ok(Self { inner })
+    }
+
+    pub fn forward(&self, batch: Batch) -> Result<(Option<Tensor>, Option<Tensor>)> {
+        let (pooled, raw) = self.inner.forward(batch)?;
+
+        // NOTE: Apply Pplx1-specific quantization to pooled embeddings
+        let pooled = pooled
+            .map(|embeddings| {
+                embeddings
+                    .tanh() // Apply tanh: [-1, 1]
+                    // NOTE: To benefit form the INT8 quantization / scaling, the `normalize`
+                    // parameter when generating embeddings should be set to `false`, otherwise the
+                    // quantization is "lost"
+                    .and_then(|t| t.affine(127.0, 0.0)) // INT8 scale: [-127, 127]
+                    .and_then(|t| t.round()) // Round to integers
+            })
+            .transpose()?;
+
+        Ok((pooled, raw))
+    }
+}
+
+impl Model for Pplx1Model {
+    fn is_padded(&self) -> bool {
+        self.inner.is_padded()
+    }
+
+    fn embed(&self, batch: Batch) -> Result<(Option<Tensor>, Option<Tensor>)> {
+        self.forward(batch)
+    }
+}

backends/candle/src/models/qwen3.rs

@@ -396,6 +396,7 @@ pub struct Qwen3Model {
     pool: Pool,
     num_attention_heads: usize,
     pad_token_id: u32,
+
     use_bidirectional_attention: bool,
 
     dtype: DType,
@@ -452,8 +453,6 @@ impl Qwen3Model {
             None
         };
 
-        let use_bidirectional_attention = config.use_bidirectional_attention.unwrap_or(false);
-
         let rotary_dim = config
             .head_dim
             .unwrap_or(config.hidden_size / config.num_attention_heads);
@@ -482,7 +481,7 @@ impl Qwen3Model {
             pool,
             pad_token_id: config.eos_token_id as u32,
             num_attention_heads: config.num_attention_heads,
-            use_bidirectional_attention,
+            use_bidirectional_attention: config.use_bidirectional_attention.unwrap_or(false),
             dtype: vb.dtype(),
             device: vb.device().clone(),
             span: tracing::span!(tracing::Level::TRACE, "model"),