candding 0.1.0

use serde::Deserialize;

use crate::error::CanddingError;
use crate::error::Result;

/// The `NewConfig` fields this family reads, as gte-multilingual-base and gte-base-en-v1.5 publish them.
#[derive(Clone, Debug, Deserialize)]
pub struct GteNewConfig {
  pub hidden_size: usize,
  pub num_hidden_layers: usize,
  pub num_attention_heads: usize,
  pub intermediate_size: usize,
  pub vocab_size: usize,
  /// Required, no default: `0` (en-v1.5) means no token-type table, `1` (mGTE) a one-row table.
  pub type_vocab_size: usize,
  #[serde(default = "default_hidden_act")]
  pub hidden_act: String,
  #[serde(default = "default_layer_norm_type")]
  pub layer_norm_type: String,
  #[serde(default = "default_layer_norm_eps")]
  pub layer_norm_eps: f64,
  #[serde(default = "default_position_embedding_type")]
  pub position_embedding_type: String,
  #[serde(default = "default_rope_theta")]
  pub rope_theta: f64,
  /// Upstream defaults to `None`, a plain unscaled rotary table.
  #[serde(default)]
  pub rope_scaling: Option<RopeScaling>,
  pub max_position_embeddings: usize,
  #[serde(default = "default_pack_qkv")]
  pub pack_qkv: bool,
  #[serde(default = "default_unpad_inputs")]
  pub unpad_inputs: bool,
  #[serde(default = "default_use_memory_efficient_attention")]
  pub use_memory_efficient_attention: bool,
}

/// `{"factor": <f64>, "type": "ntk"}`; serde strips the `r#` from `r#type`, so no rename is needed.
#[derive(Clone, Debug, Deserialize)]
pub struct RopeScaling {
  pub factor: f64,
  pub r#type: String,
}

fn default_hidden_act() -> String {
  "gelu".to_string()
}

fn default_layer_norm_type() -> String {
  "layer_norm".to_string()
}

fn default_layer_norm_eps() -> f64 {
  1e-12
}

fn default_position_embedding_type() -> String {
  "rope".to_string()
}

fn default_rope_theta() -> f64 {
  10_000.0
}

fn default_pack_qkv() -> bool {
  true
}

fn default_unpad_inputs() -> bool {
  false
}

fn default_use_memory_efficient_attention() -> bool {
  false
}

impl GteNewConfig {
  pub fn head_dim(&self) -> usize {
    self.hidden_size / self.num_attention_heads
  }

  /// `theta * factor` (NTK, applied at construction rather than past a trained length), else `rope_theta`.
  pub fn rope_theta_effective(&self) -> f64 {
    match &self.rope_scaling {
      Some(scaling) => self.rope_theta * scaling.factor,
      None => self.rope_theta,
    }
  }

  /// The uniform `factor^(-2/dim)` multiplier on every `inv_freq[i]`, else `1.0`.
  pub fn rope_frequency_scale(&self) -> f64 {
    match &self.rope_scaling {
      Some(scaling) => scaling.factor.powf(-2.0 / self.head_dim() as f64),
      None => 1.0,
    }
  }

  pub fn has_token_types(&self) -> bool {
    self.type_vocab_size > 0
  }

  /// Refuses activations, norms, position schemes, attention variants or token-type tables this port lacks.
  pub fn ensure_supported(&self) -> Result<()> {
    if self.hidden_act != "gelu" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: hidden_act {} is not supported; the port applies the erf GELU",
        self.hidden_act
      )));
    }
    if self.layer_norm_type != "layer_norm" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: layer_norm_type {} is not supported; the port always loads LayerNorm with a bias, not RMSNorm",
        self.layer_norm_type
      )));
    }
    if self.position_embedding_type != "rope" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: position_embedding_type {} is not supported; the reference's absolute branch builds a \
         position table this port has none of",
        self.position_embedding_type
      )));
    }
    if !self.pack_qkv {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: pack_qkv {} is not supported; the port reads one fused projection",
        self.pack_qkv
      )));
    }
    if self.unpad_inputs {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: unpad_inputs {} is not supported; the unpad/pad_input path is not transcribed",
        self.unpad_inputs
      )));
    }
    if self.use_memory_efficient_attention {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: use_memory_efficient_attention {} is not supported; the xformers path is not transcribed",
        self.use_memory_efficient_attention
      )));
    }
    if let Some(scaling) = &self.rope_scaling
      && scaling.r#type != "ntk"
    {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: rope_scaling.type {} is not supported; only \"ntk\" is implemented, the reference's linear \
         and dynamic branches are commented out",
        scaling.r#type
      )));
    }
    if self.type_vocab_size > 1 {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "gte_new: type_vocab_size {} is not supported; the reference zeroes ids below 2, so a real second \
         segment would change the forward",
        self.type_vocab_size
      )));
    }
    Ok(())
  }
}

#[cfg(test)]
mod tests;