candding 0.1.0

use serde::Deserialize;

use super::bias::NUM_BUCKETS;
use crate::error::CanddingError;
use crate::error::Result;
use crate::layers::activation::Activation;

/// transformers hard-codes MPNet's padding index to 1 rather than reading `pad_token_id`; positions start at 2.
pub(super) const PADDING_IDX: u32 = 1;
pub(super) const POSITION_OFFSET: u32 = PADDING_IDX + 1;

/// Not flattened over `BertConfig`: its `type_vocab_size` default would invent a token-type table MPNet lacks.
#[derive(Clone, Debug, Deserialize)]
pub struct MPNetConfig {
  pub vocab_size: usize,
  pub hidden_size: usize,
  pub num_hidden_layers: usize,
  pub num_attention_heads: usize,
  pub intermediate_size: usize,
  #[serde(default = "default_act")]
  pub hidden_act: Activation,
  pub max_position_embeddings: usize,
  #[serde(default = "default_eps")]
  pub layer_norm_eps: f64,
  #[serde(default = "default_num_buckets")]
  pub relative_attention_num_buckets: usize,
  /// Read only so a test can record that it equals `PADDING_IDX`; the numbering does not consult it.
  #[serde(default = "default_pad_token_id")]
  pub pad_token_id: usize,
}

fn default_act() -> Activation {
  Activation::Gelu
}

/// transformers' default; every catalog checkpoint declares 1e-5.
fn default_eps() -> f64 {
  1e-12
}

fn default_num_buckets() -> usize {
  NUM_BUCKETS
}

fn default_pad_token_id() -> usize {
  PADDING_IDX as usize
}

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

  /// The longest sequence the position table can index: its first two rows are never positions, so 514 rows give 512.
  pub fn sequence_limit(&self) -> usize {
    self
      .max_position_embeddings
      .saturating_sub(POSITION_OFFSET as usize)
  }

  /// The reference buckets with the default 32 whatever the table size, so any other value would silently diverge.
  pub fn ensure_supported(&self) -> Result<()> {
    if self.relative_attention_num_buckets != NUM_BUCKETS {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "mpnet with {} relative attention buckets",
        self.relative_attention_num_buckets
      )));
    }
    Ok(())
  }
}

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

  /// `sentence-transformers/all-mpnet-base-v2`'s `config.json`: no `type_vocab_size` and no `position_embedding_type`.
  const ALL_MPNET_BASE_V2: &str = r#"{
    "architectures": ["MPNetForMaskedLM"],
    "attention_probs_dropout_prob": 0.1,
    "bos_token_id": 0,
    "eos_token_id": 2,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-05,
    "max_position_embeddings": 514,
    "model_type": "mpnet",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "pad_token_id": 1,
    "relative_attention_num_buckets": 32,
    "transformers_version": "4.8.2",
    "vocab_size": 30527
  }"#;

  #[test]
  fn reads_the_published_config() {
    let config: MPNetConfig = serde_json::from_str(ALL_MPNET_BASE_V2).unwrap();
    assert_eq!(config.vocab_size, 30527);
    assert_eq!(config.hidden_size, 768);
    assert_eq!(config.num_hidden_layers, 12);
    assert_eq!(config.num_attention_heads, 12);
    assert_eq!(config.intermediate_size, 3072);
    assert_eq!(config.hidden_act, Activation::Gelu);
    assert_eq!(config.max_position_embeddings, 514);
    assert_eq!(config.layer_norm_eps, 1e-5);
    assert_eq!(config.relative_attention_num_buckets, 32);
    assert_eq!(config.head_dim(), 64);
  }

  /// `padding_idx + 1` rows precede the first real position, so 514 rows carry 512 positions.
  #[test]
  fn the_sequence_limit_drops_the_position_offset() {
    let config: MPNetConfig = serde_json::from_str(ALL_MPNET_BASE_V2).unwrap();
    assert_eq!(config.sequence_limit(), 512);
    assert_eq!(600usize.min(config.sequence_limit()), 512);
    assert_eq!(200usize.min(config.sequence_limit()), 200);
  }

  /// `MPNetEmbeddings.__init__` sets `padding_idx = 1` as a literal; this records that the config agrees.
  #[test]
  fn the_pad_token_id_agrees_with_the_literal_the_numbering_uses() {
    let config: MPNetConfig = serde_json::from_str(ALL_MPNET_BASE_V2).unwrap();
    assert_eq!(config.pad_token_id as u32, PADDING_IDX);
    assert_eq!(POSITION_OFFSET, PADDING_IDX + 1);
  }

  /// MPNet has no token-type table, so the config is standalone rather than flattening `BertConfig`'s default of 2.
  #[test]
  fn a_token_type_field_is_ignored_rather_than_honoured() {
    assert!(!ALL_MPNET_BASE_V2.contains("type_vocab_size"));
    let with_field = ALL_MPNET_BASE_V2.replace(
      "\"vocab_size\": 30527",
      "\"type_vocab_size\": 2,\n    \"vocab_size\": 30527",
    );
    let config: MPNetConfig = serde_json::from_str(&with_field).unwrap();
    assert_eq!(config.vocab_size, 30527);
  }

  /// `transformers`' `MPNetConfig` defaults; every catalog checkpoint overrides `layer_norm_eps` with 1e-5.
  #[test]
  fn the_optional_fields_default_to_the_reference_config_values() {
    let stripped = ALL_MPNET_BASE_V2
      .replace("\"hidden_act\": \"gelu\",", "")
      .replace("\"layer_norm_eps\": 1e-05,", "")
      .replace("\"pad_token_id\": 1,", "")
      .replace("\"relative_attention_num_buckets\": 32,", "");
    let config: MPNetConfig = serde_json::from_str(&stripped).unwrap();
    assert_eq!(config.hidden_act, Activation::Gelu);
    assert_eq!(config.layer_norm_eps, 1e-12);
    assert_eq!(config.relative_attention_num_buckets, 32);
    assert_eq!(config.pad_token_id, 1);
  }

  /// The reference buckets with a literal 32 whatever the config says, so another bucket count is refused at load.
  #[test]
  fn only_thirty_two_buckets_are_supported() {
    let supported: MPNetConfig = serde_json::from_str(ALL_MPNET_BASE_V2).unwrap();
    supported.ensure_supported().unwrap();
    let wide: MPNetConfig = serde_json::from_str(&ALL_MPNET_BASE_V2.replace(
      "\"relative_attention_num_buckets\": 32",
      "\"relative_attention_num_buckets\": 64",
    ))
    .unwrap();
    let err = wide.ensure_supported().unwrap_err();
    assert!(
      matches!(err, CanddingError::UnsupportedArchitecture(ref m) if m.contains("64")),
      "{err}"
    );
  }
}