Fix max_position_embeddings handling in NomicBertConfig and FlashNomicBertConfig (#876)

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

Author: alvarobartt

backends/candle/src/models/flash_nomic.rs

@@ -196,7 +196,7 @@ pub struct FlashNomicBertModel {
     pool: Pool,
     pub device: Device,
 
-    max_trained_positions: u32,
+    max_position_embeddings: u32,
     rotary_cache: (Tensor, Tensor),
     scaled_rotary_cache: Option<(Tensor, Tensor)>,
 
@@ -233,14 +233,17 @@ impl FlashNomicBertModel {
         let embeddings = NomicBertEmbeddings::load(vb.clone(), config)?;
         let encoder = NomicBertEncoder::load(vb.pp("encoder"), config)?;
 
+        let max_position_embeddings = config
+            .max_position_embeddings
+            .unwrap_or(config.max_trained_positions.unwrap_or(2048));
+
         let rotary_dim = encoder.layers[0].attention.attention_head_size;
         let inv_freqs = get_inv_freqs(rotary_dim, config.rotary_emb_base, vb.device(), None)?;
         let rotary_cache = get_cos_sin(config.n_positions, &inv_freqs, vb.dtype(), false)?;
 
         let scaled_rotary_cache = if let Some(scaling_factor) = config.rotary_scaling_factor {
             let new_base = (config.rotary_emb_base
-                * ((scaling_factor * config.n_positions as f32
-                    / config.max_trained_positions as f32)
+                * ((scaling_factor * config.n_positions as f32 / max_position_embeddings as f32)
                     - (scaling_factor - 1.0)))
                 .powi((rotary_dim as f32 / (rotary_dim as f32 - 2.0)) as i32);
             let inv_freqs = get_inv_freqs(rotary_dim, new_base, vb.device(), None)?;
@@ -258,7 +261,7 @@ impl FlashNomicBertModel {
             embeddings,
             encoder,
             pool,
-            max_trained_positions: config.max_trained_positions as u32,
+            max_position_embeddings: max_position_embeddings as u32,
             rotary_cache,
             scaled_rotary_cache,
             device: vb.device().clone(),
@@ -283,7 +286,7 @@ impl FlashNomicBertModel {
         )?;
 
         let (cos, sin) = if self.scaled_rotary_cache.is_some()
-            && batch.max_length > self.max_trained_positions
+            && batch.max_length > self.max_position_embeddings
         {
             let cos = index_select(
                 &self.scaled_rotary_cache.as_ref().unwrap().0,

backends/candle/src/models/nomic.rs

@@ -18,8 +18,13 @@ pub struct NomicConfig {
     pub mlp_fc1_bias: bool,
     pub mlp_fc2_bias: bool,
     pub rotary_scaling_factor: Option<f32>,
-    #[serde(default = "default_max_trained_positions")]
-    pub max_trained_positions: usize,
+
+    // NOTE: `max_trained_positions` is specific for NomicBERT when it required custom code, but
+    // since Transformers v5 it's no longer required, and it now defines `max_position_embeddings`
+    // in the `config.json` instead. Not included as an `alias` since both can be present at the
+    // same time, see https://huggingface.co/nomic-ai/nomic-embed-text-v1.5/blob/e9b6763023c676ca8431644204f50c2b100d9aab/config.json#L33-L34
+    pub max_trained_positions: Option<usize>,
+    pub max_position_embeddings: Option<usize>,
 
     pub moe_every_n_layers: Option<usize>,
     pub moe_normalize_expert_weights: Option<bool>,
@@ -39,10 +44,6 @@ pub struct NomicConfig {
     pub layer_norm_epsilon: f32,
 }
 
-fn default_max_trained_positions() -> usize {
-    2048
-}
-
 impl NomicConfig {
     // For now, we only support these parameters
     pub fn valid(&self) -> bool {
@@ -668,7 +669,7 @@ pub struct NomicBertModel {
     dtype: DType,
 
     rotary_dim: usize,
-    max_trained_positions: u32,
+    max_position_embeddings: u32,
     rotary_cache: (Tensor, Tensor),
     scaled_rotary_cache: Option<(Tensor, Tensor)>,
 
@@ -702,15 +703,18 @@ impl NomicBertModel {
         let embeddings = NomicBertEmbeddings::load(vb.clone(), config)?;
         let encoder = NomicBertEncoder::load(vb.pp("encoder"), config)?;
 
+        let max_position_embeddings = config
+            .max_position_embeddings
+            .unwrap_or(config.max_trained_positions.unwrap_or(2048));
+
         let rotary_dim = encoder.layers[0].attention.attention_head_size;
         let inv_freqs_tensor =
             get_inv_freqs(rotary_dim, config.rotary_emb_base, vb.device(), None)?;
         let rotary_cache = get_cos_sin(config.n_positions, &inv_freqs_tensor, vb.dtype(), true)?;
 
         let scaled_rotary_cache = if let Some(scaling_factor) = config.rotary_scaling_factor {
             let new_base = (config.rotary_emb_base
-                * ((scaling_factor * config.n_positions as f32
-                    / config.max_trained_positions as f32)
+                * ((scaling_factor * config.n_positions as f32 / max_position_embeddings as f32)
                     - (scaling_factor - 1.0)))
                 .powi((rotary_dim as f32 / (rotary_dim as f32 - 2.0)) as i32);
             let inv_freqs_tensor = get_inv_freqs(rotary_dim, new_base, vb.device(), None)?;
@@ -729,7 +733,7 @@ impl NomicBertModel {
             encoder,
             pool,
             rotary_dim,
-            max_trained_positions: config.max_trained_positions as u32,
+            max_position_embeddings: max_position_embeddings as u32,
             rotary_cache,
             scaled_rotary_cache,
             num_attention_heads: config.n_head,
@@ -855,7 +859,7 @@ impl NomicBertModel {
             Tensor::from_vec(input_lengths, (batch_size, 1), &self.device)?.to_dtype(self.dtype)?;
 
         let (cos, sin) = if self.scaled_rotary_cache.is_some()
-            && batch.max_length > self.max_trained_positions
+            && batch.max_length > self.max_position_embeddings
         {
             let cos = self
                 .scaled_rotary_cache

router/src/lib.rs

@@ -202,11 +202,19 @@ pub async fn run(
         }
     }
 
+    let max_position_embeddings = match config.max_position_embeddings {
+        Some(max_position_embeddings) => max_position_embeddings,
+        None => match config.max_trained_positions {
+            Some(max_trained_positions) => max_trained_positions,
+            None => anyhow::bail!("At least any of `max_position_embeddings` or `max_trained_positions` (only applies for NomicBERT), in that order of priority, should be defined in `config.json`."),
+        },
+    };
+
     let base_input_length = match st_config {
         Some(config) => config.max_seq_length,
         None => {
             tracing::warn!("Could not find a Sentence Transformers config");
-            config.max_position_embeddings - position_offset
+            max_position_embeddings - position_offset
         }
     };
 
@@ -461,8 +469,12 @@ fn get_backend_model_type(
 pub struct ModelConfig {
     pub architectures: Vec<String>,
     pub model_type: String,
-    #[serde(alias = "n_positions")]
-    pub max_position_embeddings: usize,
+    // NOTE: `max_trained_positions` is specific for NomicBERT when it required custom code, but
+    // since Transformers v5 it's no longer required, and it now defines `max_position_embeddings`
+    // in the `config.json` instead. Not included as an `alias` since both can be present at the
+    // same time, see https://huggingface.co/nomic-ai/nomic-embed-text-v1.5/blob/e9b6763023c676ca8431644204f50c2b100d9aab/config.json#L33-L34
+    pub max_trained_positions: Option<usize>,
+    pub max_position_embeddings: Option<usize>,
     #[serde(default)]
     pub pad_token_id: Option<usize>,
     pub id2label: Option<HashMap<String, String>>,