candding 0.1.0

//! A zero-layer encoder with a zeroed pooler: `tanh(0) == 0`, so every score is exactly `classifier.bias`.
//! Configs are literal JSON, since neither config type derives `Serialize`.

use std::collections::HashMap;

use candle_core::Device;

use super::*;

fn bert_config_json(hidden: usize) -> String {
  format!(
    r#"{{"vocab_size":64,"hidden_size":{hidden},"num_hidden_layers":0,"num_attention_heads":1,
    "intermediate_size":1,"max_position_embeddings":32,"type_vocab_size":2,"layer_norm_eps":1e-12}}"#
  )
}

fn jina_config_json(hidden: usize) -> String {
  format!(
    r#"{{"vocab_size":64,"hidden_size":{hidden},"num_hidden_layers":0,"num_attention_heads":1,
    "intermediate_size":1,"max_position_embeddings":32,"type_vocab_size":2,"layer_norm_eps":1e-12,
    "position_embedding_type":"alibi","feed_forward_type":"geglu"}}"#
  )
}

fn embeddings_tensors(hidden: usize, prefix: &str, position: bool) -> HashMap<String, Tensor> {
  let mut tensors = HashMap::new();
  let mut put = |name: &str, tensor: Tensor| {
    tensors.insert(format!("{prefix}{name}"), tensor);
  };
  put(
    "embeddings.word_embeddings.weight",
    Tensor::zeros((64, hidden), DType::F32, &Device::Cpu).unwrap(),
  );
  if position {
    put(
      "embeddings.position_embeddings.weight",
      Tensor::zeros((32, hidden), DType::F32, &Device::Cpu).unwrap(),
    );
  }
  put(
    "embeddings.token_type_embeddings.weight",
    Tensor::zeros((2, hidden), DType::F32, &Device::Cpu).unwrap(),
  );
  put(
    "embeddings.LayerNorm.weight",
    Tensor::ones(hidden, DType::F32, &Device::Cpu).unwrap(),
  );
  put(
    "embeddings.LayerNorm.bias",
    Tensor::zeros(hidden, DType::F32, &Device::Cpu).unwrap(),
  );
  tensors
}

fn head_tensors(hidden: usize, classifier_bias: f32) -> HashMap<String, Tensor> {
  let mut tensors = HashMap::new();
  tensors.insert(
    "bert.pooler.dense.weight".to_string(),
    Tensor::zeros((hidden, hidden), DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "bert.pooler.dense.bias".to_string(),
    Tensor::zeros(hidden, DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "classifier.weight".to_string(),
    Tensor::zeros((1, hidden), DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "classifier.bias".to_string(),
    Tensor::new(&[classifier_bias], &Device::Cpu).unwrap(),
  );
  tensors
}

/// `token_type_ids` carry a real `1` run past the first `[SEP]`: a genuine pair.
fn pair_input() -> (Tensor, Tensor, Tensor, Vec<usize>) {
  let input_ids = Tensor::new(&[[1u32, 2, 3, 0], [1u32, 4, 5, 6]], &Device::Cpu).unwrap();
  let attention_mask = Tensor::new(&[[1u32, 1, 1, 0], [1u32, 1, 1, 1]], &Device::Cpu).unwrap();
  let token_type_ids = Tensor::new(&[[0u32, 0, 1, 0], [0u32, 1, 1, 1]], &Device::Cpu).unwrap();
  (input_ids, attention_mask, token_type_ids, vec![3, 4])
}

fn scores_for(model: &BertPoolerRerank) -> Vec<f32> {
  let (input_ids, attention_mask, token_type_ids, lengths) = pair_input();
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  model
    .forward(&input)
    .unwrap()
    .into_iter()
    .map(|s| s.value())
    .collect()
}

#[test]
fn bert_family_loads_and_scores_every_row_as_the_classifier_bias_under_identity() {
  let mut tensors = embeddings_tensors(4, "bert.", true);
  tensors.extend(head_tensors(4, 8.845_856));
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let model = BertPoolerRerank::load(
    "bert",
    &bert_config_json(4),
    vb,
    ScoreActivation::Identity,
    None,
  )
  .unwrap();
  let scores = scores_for(&model);
  assert_eq!(scores.len(), 2);
  for score in scores {
    assert!((score - 8.845_856).abs() < 1e-5, "{score}");
  }
}

/// `forward` must read `score_activation`, since a raw logit and its sigmoid are both plausible scores.
#[test]
fn sigmoid_activation_squashes_the_same_raw_logit() {
  let mut tensors = embeddings_tensors(4, "bert.", true);
  tensors.extend(head_tensors(4, 0.0));
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let model = BertPoolerRerank::load(
    "bert",
    &bert_config_json(4),
    vb,
    ScoreActivation::Sigmoid,
    None,
  )
  .unwrap();
  for score in scores_for(&model) {
    assert!((score - 0.5).abs() < 1e-5, "{score}");
  }
}

#[test]
fn a_hidden_size_that_disagrees_with_the_real_tensor_shape_is_refused() {
  let mut tensors = embeddings_tensors(4, "bert.", true);
  tensors.extend(head_tensors(4, 1.0));
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  // hidden_size 6 in the config disagrees with the real, 4-wide tensors above.
  match BertPoolerRerank::load(
    "bert",
    &bert_config_json(6),
    vb,
    ScoreActivation::Identity,
    None,
  ) {
    Ok(_) => panic!("a hidden size that disagrees with the real tensors must be refused"),
    Err(err) => assert!(matches!(err, CanddingError::Candle(_)), "{err}"),
  }
}

/// `JinaBertModel::load` has no prefix probe, so the `jina_bert` arm strips `bert.` itself.
#[test]
fn jina_bert_family_strips_its_own_prefix_and_scores_every_row() {
  let mut tensors = embeddings_tensors(4, "bert.", false);
  tensors.extend(head_tensors(4, -11.245_562));
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let model = BertPoolerRerank::load(
    "jina_bert",
    &jina_config_json(4),
    vb,
    ScoreActivation::Identity,
    None,
  )
  .unwrap();
  let scores = scores_for(&model);
  assert_eq!(scores.len(), 2);
  for score in scores {
    assert!((score - -11.245_562).abs() < 1e-4, "{score}");
  }
}

#[test]
fn an_unrecognized_model_type_is_refused_by_name() {
  let vb = VarBuilder::from_tensors(HashMap::new(), DType::F32, &Device::Cpu);
  match BertPoolerRerank::load("xlm-roberta", "{}", vb, ScoreActivation::Identity, None) {
    Ok(_) => panic!("an unrecognized model_type must be refused"),
    Err(CanddingError::UnsupportedArchitecture(got)) => assert_eq!(got, "xlm-roberta"),
    Err(other) => panic!("unexpected error: {other}"),
  }
}