candding 0.1.0

use serde::Deserialize;

use crate::error::CanddingError;
use crate::error::Result;
use crate::models::bert::config::BertConfig;

/// `BertConfig` flattened in, plus the two fields the position scheme needs.
#[derive(Clone, Debug, Deserialize)]
pub struct XlmRobertaConfig {
  #[serde(flatten)]
  pub bert: BertConfig,
  #[serde(default = "default_pad_token_id")]
  pub pad_token_id: usize,
  #[serde(default = "default_position_embedding_type")]
  pub position_embedding_type: String,
}

fn default_pad_token_id() -> usize {
  1
}

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

impl XlmRobertaConfig {
  pub fn head_dim(&self) -> usize {
    self.bert.head_dim()
  }

  /// The first real token's position id: the rows up to and including `padding_idx` are reserved.
  pub fn position_offset(&self) -> u32 {
    self.pad_token_id as u32 + 1
  }

  /// The longest sequence the position table can index: 514 rows carry 512 positions, 8194 carry 8192.
  pub fn sequence_limit(&self) -> usize {
    self
      .bert
      .max_position_embeddings
      .saturating_sub(self.pad_token_id + 1)
  }

  /// Refuses `relative_key` and `relative_key_query`, which no catalog checkpoint uses and candding lacks.
  pub fn ensure_supported(&self) -> Result<()> {
    if self.position_embedding_type != "absolute" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "xlm-roberta with {} position embeddings",
        self.position_embedding_type
      )));
    }
    Ok(())
  }
}

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

  /// `intfloat/multilingual-e5-base`'s `config.json`, byte for byte apart from the removed `_name_or_path`.
  const E5_BASE: &str = r#"{
    "architectures": ["XLMRobertaModel"],
    "attention_probs_dropout_prob": 0.1,
    "bos_token_id": 0,
    "classifier_dropout": null,
    "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": "xlm-roberta",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "output_past": true,
    "pad_token_id": 1,
    "position_embedding_type": "absolute",
    "torch_dtype": "float32",
    "transformers_version": "4.29.0.dev0",
    "type_vocab_size": 1,
    "use_cache": true,
    "vocab_size": 250002
  }"#;

  /// `BAAI/bge-m3`'s `config.json`: the same body at 24 layers, 1024 hidden and 8194 positions.
  const BGE_M3: &str = r#"{
    "architectures": ["XLMRobertaModel"],
    "attention_probs_dropout_prob": 0.1,
    "bos_token_id": 0,
    "classifier_dropout": null,
    "eos_token_id": 2,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 1024,
    "initializer_range": 0.02,
    "intermediate_size": 4096,
    "layer_norm_eps": 1e-05,
    "max_position_embeddings": 8194,
    "model_type": "xlm-roberta",
    "num_attention_heads": 16,
    "num_hidden_layers": 24,
    "output_past": true,
    "pad_token_id": 1,
    "position_embedding_type": "absolute",
    "torch_dtype": "float32",
    "transformers_version": "4.33.0",
    "type_vocab_size": 1,
    "use_cache": true,
    "vocab_size": 250002
  }"#;

  #[test]
  fn reads_the_shared_bert_fields_through_the_flatten() {
    let config: XlmRobertaConfig = serde_json::from_str(E5_BASE).unwrap();
    assert_eq!(config.bert.vocab_size, 250002);
    assert_eq!(config.bert.hidden_size, 768);
    assert_eq!(config.bert.num_hidden_layers, 12);
    assert_eq!(config.bert.num_attention_heads, 12);
    assert_eq!(config.bert.intermediate_size, 3072);
    assert_eq!(config.bert.hidden_act, Activation::Gelu);
    assert_eq!(config.bert.type_vocab_size, 1);
    assert_eq!(config.bert.layer_norm_eps, 1e-5);
    assert_eq!(config.head_dim(), 64);
  }

  /// `pad_token_id + 1` rows precede the first real position: 514 rows carry 512 positions and 8194 carry 8192.
  #[test]
  fn the_sequence_limit_drops_the_position_offset() {
    let base: XlmRobertaConfig = serde_json::from_str(E5_BASE).unwrap();
    assert_eq!(base.pad_token_id, 1);
    assert_eq!(base.position_offset(), 2);
    assert_eq!(base.bert.max_position_embeddings, 514);
    assert_eq!(base.sequence_limit(), 512);
    let m3: XlmRobertaConfig = serde_json::from_str(BGE_M3).unwrap();
    assert_eq!(m3.bert.hidden_size, 1024);
    assert_eq!(m3.bert.num_hidden_layers, 24);
    assert_eq!(m3.head_dim(), 64);
    assert_eq!(m3.sequence_limit(), 8192);
  }

  #[test]
  fn absolute_positions_are_the_only_supported_scheme() {
    let absolute: XlmRobertaConfig = serde_json::from_str(E5_BASE).unwrap();
    absolute.ensure_supported().unwrap();
    let relative: XlmRobertaConfig = serde_json::from_str(&E5_BASE.replace(
      "\"position_embedding_type\": \"absolute\"",
      "\"position_embedding_type\": \"relative_key\"",
    ))
    .unwrap();
    let err = relative.ensure_supported().unwrap_err();
    assert!(
      matches!(err, CanddingError::UnsupportedArchitecture(ref m) if m.contains("relative_key")),
      "{err}"
    );
  }

  #[test]
  fn the_two_own_fields_have_defaults() {
    let stripped = E5_BASE
      .replace("\"pad_token_id\": 1,", "")
      .replace("\"position_embedding_type\": \"absolute\",", "");
    let config: XlmRobertaConfig = serde_json::from_str(&stripped).unwrap();
    assert_eq!(config.pad_token_id, 1);
    config.ensure_supported().unwrap();
  }
}