candding 0.1.0

//! A zero-layer encoder with a zeroed `dense`: `tanh(0) == 0`, so every score is exactly `out_proj.bias`.

use std::collections::HashMap;

use candle_core::Device;

use super::*;

fn 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":1,"layer_norm_eps":1e-12,
    "pad_token_id":1,"position_embedding_type":"absolute"}}"#
  )
}

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

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

/// XLM-RoBERTa marks a pair only by the double separator, so `token_type_ids` are all zero.
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, 0, 0], [0u32, 0, 0, 0]], &Device::Cpu).unwrap();
  (input_ids, attention_mask, token_type_ids, vec![3, 4])
}

fn scores_for(model: &RobertaClassificationRerank) -> 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 loads_and_scores_every_row_as_the_out_proj_bias_under_identity() {
  let mut tensors = embeddings_tensors(4);
  tensors.extend(head_tensors(4, -5.652_344));
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let model = RobertaClassificationRerank::load(
    "xlm-roberta",
    &config_json(4),
    vb,
    ScoreActivation::Identity,
    None,
  )
  .unwrap();
  let scores = scores_for(&model);
  assert_eq!(scores.len(), 2);
  for score in scores {
    assert!((score - -5.652_344).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);
  tensors.extend(head_tensors(4, 0.0));
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let model = RobertaClassificationRerank::load(
    "xlm-roberta",
    &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);
  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 RobertaClassificationRerank::load(
    "xlm-roberta",
    &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}"),
  }
}

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

mod flash;