candding 0.1.0

use serde::Deserialize;

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

/// Not flattened over `BertConfig`, which would hide `position_embedding_type` and `feed_forward_type`.
#[derive(Clone, Debug, Deserialize)]
pub struct JinaBertConfig {
  pub vocab_size: usize,
  pub hidden_size: usize,
  pub num_hidden_layers: usize,
  pub num_attention_heads: usize,
  pub intermediate_size: usize,
  /// 2 in every checkpoint; row 0 reaches every single-sequence input.
  #[serde(default = "default_type_vocab_size")]
  pub type_vocab_size: usize,
  #[serde(default = "default_eps")]
  pub layer_norm_eps: f64,
  /// The trained range and memory cap, not a table height: ALiBi indexes nothing.
  pub max_position_embeddings: usize,
  /// No default: a config without it is not a JinaBERT.
  pub position_embedding_type: String,
  /// Picks the gated block's activation, not `hidden_act`; upstream's default `"original"` is the non-gated block.
  #[serde(default = "default_feed_forward_type")]
  pub feed_forward_type: String,
}

fn default_type_vocab_size() -> usize {
  2
}

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

fn default_feed_forward_type() -> String {
  "original".to_string()
}

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

  /// `absolute` would leave a position table unread and `reglu` would gate with ReLU, both silently wrong.
  pub fn ensure_supported(&self) -> Result<()> {
    if self.position_embedding_type != "alibi" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "jina_bert with position_embedding_type `{}`",
        self.position_embedding_type
      )));
    }
    if self.feed_forward_type != "geglu" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "jina_bert with feed_forward_type `{}`",
        self.feed_forward_type
      )));
    }
    Ok(())
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use crate::error::CanddingError;

  /// `auto_map` is kept because candding must be seen to ignore it: following it would execute remote code.
  const BASE_EN: &str = r#"{
    "architectures": ["JinaBertForMaskedLM"],
    "attention_probs_dropout_prob": 0.0,
    "auto_map": {
      "AutoConfig": "jinaai/jina-bert-implementation--configuration_bert.JinaBertConfig",
      "AutoModel": "jinaai/jina-bert-implementation--modeling_bert.JinaBertModel",
      "AutoModelForMaskedLM": "jinaai/jina-bert-implementation--modeling_bert.JinaBertForMaskedLM"
    },
    "emb_pooler": "mean",
    "feed_forward_type": "geglu",
    "gradient_checkpointing": false,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-12,
    "max_position_embeddings": 8192,
    "model_max_length": 8192,
    "model_type": "bert",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "pad_token_id": 0,
    "position_embedding_type": "alibi",
    "torch_dtype": "float32",
    "type_vocab_size": 2,
    "vocab_size": 30528
  }"#;

  /// The other five differ only in five sizes and `auto_map`, so they are built by substitution.
  fn config_json(
    hidden: usize,
    layers: usize,
    heads: usize,
    intermediate: usize,
    vocab: usize,
  ) -> String {
    BASE_EN
      .replace(
        "\"hidden_size\": 768",
        &format!("\"hidden_size\": {hidden}"),
      )
      .replace(
        "\"num_hidden_layers\": 12",
        &format!("\"num_hidden_layers\": {layers}"),
      )
      .replace(
        "\"num_attention_heads\": 12",
        &format!("\"num_attention_heads\": {heads}"),
      )
      .replace(
        "\"intermediate_size\": 3072",
        &format!("\"intermediate_size\": {intermediate}"),
      )
      .replace("\"vocab_size\": 30528", &format!("\"vocab_size\": {vocab}"))
  }

  /// The head size is 64 for all six, inside `SDPA_HEAD_DIMS` and below `SDPA_F32_MAX_HEAD_DIM`, so Metal's kernel runs in every dtype.
  #[test]
  fn reads_all_six_published_configs() {
    let cases = [
      ("small-en", 512usize, 4usize, 8usize, 2048usize, 30528usize),
      ("base-en", 768, 12, 12, 3072, 30528),
      ("base-de", 768, 12, 12, 3072, 61056),
      ("base-es", 768, 12, 12, 3072, 61056),
      ("base-zh", 768, 12, 12, 3072, 61056),
      ("base-code", 768, 12, 12, 3072, 61056),
    ];
    for (name, hidden, layers, heads, intermediate, vocab) in cases {
      let config: JinaBertConfig =
        serde_json::from_str(&config_json(hidden, layers, heads, intermediate, vocab)).unwrap();
      assert_eq!(config.hidden_size, hidden, "{name}");
      assert_eq!(config.num_hidden_layers, layers, "{name}");
      assert_eq!(config.num_attention_heads, heads, "{name}");
      assert_eq!(config.intermediate_size, intermediate, "{name}");
      assert_eq!(config.vocab_size, vocab, "{name}");
      assert_eq!(config.head_dim(), 64, "{name}");
      assert_eq!(config.type_vocab_size, 2, "{name}");
      assert_eq!(config.layer_norm_eps, 1e-12, "{name}");
      assert_eq!(config.max_position_embeddings, 8192, "{name}");
      assert_eq!(config.position_embedding_type, "alibi", "{name}");
      assert_eq!(config.feed_forward_type, "geglu", "{name}");
      config.ensure_supported().unwrap();
    }
  }

  #[test]
  fn the_head_dim_divides_by_the_head_count() {
    let config: JinaBertConfig =
      serde_json::from_str(&config_json(768, 3, 12, 3072, 30528)).unwrap();
    assert_eq!(config.head_dim(), 64);
  }

  /// `feed_forward_type` defaults to `"original"`, the non-gated feed-forward `ensure_supported` refuses.
  #[test]
  fn the_optional_fields_default_to_the_reference_config_values() {
    let stripped = BASE_EN
      .replace("\"feed_forward_type\": \"geglu\",", "")
      .replace("\"layer_norm_eps\": 1e-12,", "")
      .replace("\"type_vocab_size\": 2,", "");
    let config: JinaBertConfig = serde_json::from_str(&stripped).unwrap();
    assert_eq!(config.type_vocab_size, 2);
    assert_eq!(config.layer_norm_eps, 1e-12);
    assert_eq!(config.feed_forward_type, "original");
    let err = config.ensure_supported().unwrap_err();
    assert!(
      matches!(err, CanddingError::UnsupportedArchitecture(ref m) if m.contains("original")),
      "{err}"
    );
  }

  /// Without `position_embedding_type` this is not a JinaBERT, so a missing field fails to deserialize.
  #[test]
  fn a_missing_position_scheme_is_a_deserialization_error() {
    let stripped = BASE_EN.replace("\"position_embedding_type\": \"alibi\",", "");
    assert!(serde_json::from_str::<JinaBertConfig>(&stripped).is_err());
  }

  /// Both would otherwise load the same tensors and produce plausible wrong numbers.
  #[test]
  fn an_absolute_scheme_and_a_reglu_feed_forward_are_refused() {
    let absolute: JinaBertConfig =
      serde_json::from_str(&BASE_EN.replace("\"alibi\"", "\"absolute\"")).unwrap();
    let err = absolute.ensure_supported().unwrap_err();
    assert!(
      matches!(err, CanddingError::UnsupportedArchitecture(ref m) if m.contains("absolute")),
      "{err}"
    );
    let reglu: JinaBertConfig =
      serde_json::from_str(&BASE_EN.replace("\"geglu\"", "\"reglu\"")).unwrap();
    let err = reglu.ensure_supported().unwrap_err();
    assert!(
      matches!(err, CanddingError::UnsupportedArchitecture(ref m) if m.contains("reglu")),
      "{err}"
    );
  }
}