candding 0.1.0

use candle_core::DType;
use candle_core::Device;
use candle_core::Shape;
use candle_nn::Init;
use candle_nn::VarMap;

use super::*;

mod shapes;

fn tiny_config() -> JinaBertConfig {
  JinaBertConfig {
    vocab_size: 32,
    hidden_size: 16,
    num_hidden_layers: 2,
    num_attention_heads: 4,
    intermediate_size: 32,
    type_vocab_size: 2,
    layer_norm_eps: 1e-12,
    max_position_embeddings: 8192,
    position_embedding_type: "alibi".to_string(),
    feed_forward_type: "geglu".to_string(),
  }
}

/// Seeds one tensor before the `VarBuilder` is built, the only way a fixture decides what the load finds.
fn touch(varmap: &VarMap, name: &str, shape: impl Into<Shape>, init: Init) {
  varmap
    .get(shape, name, init, DType::F32, &Device::Cpu)
    .unwrap();
}

/// A fresh `VarMap` probes as the base shape, so a variant model seeds its marker first.
/// The fused feed-forward is seeded explicitly: `VarBuilder::get` zero-initializes, which would leave no feed-forward to test.
fn tiny_model_with(varmap: &VarMap, qk_norm: bool, split_norms: bool) -> JinaBertModel {
  let config = tiny_config();
  if qk_norm {
    touch(
      varmap,
      "encoder.layer.0.attention.self.layer_norm_q.weight",
      config.hidden_size,
      Init::Const(1.0),
    );
  }
  if split_norms {
    touch(
      varmap,
      "encoder.layer.0.layer_norm_1.weight",
      config.hidden_size,
      Init::Const(1.0),
    );
  }
  let fused = match split_norms {
    true => "mlp.up_gated_layer",
    false => "mlp.gated_layers",
  };
  for layer in 0..config.num_hidden_layers {
    touch(
      varmap,
      &format!("encoder.layer.{layer}.{fused}.weight"),
      (2 * config.intermediate_size, config.hidden_size),
      candle_nn::init::DEFAULT_KAIMING_NORMAL,
    );
  }
  let vb = VarBuilder::from_varmap(varmap, DType::F32, &Device::Cpu);
  JinaBertModel::load(config, vb, AlibiSlopeVariant::Canonical).unwrap()
}

/// `VarMap::set_one` writes into the storage the loaded model holds, so a test can rewrite one tensor in place.
fn tiny_model() -> (VarMap, JinaBertModel) {
  let varmap = VarMap::new();
  let model = tiny_model_with(&varmap, false, false);
  (varmap, model)
}

/// The padding id is 0 for all six checkpoints — every one declares `"pad_token_id": 0` — so the rows pad with 0.
fn run(model: &JinaBertModel, rows: &[Vec<u32>]) -> Vec<Vec<Vec<f32>>> {
  let seq = rows.iter().map(Vec::len).max().unwrap();
  let mut ids = Vec::new();
  let mut mask = Vec::new();
  for row in rows {
    ids.extend(
      row
        .iter()
        .copied()
        .chain(std::iter::repeat_n(0, seq - row.len())),
    );
    mask
      .extend(std::iter::repeat_n(1u32, row.len()).chain(std::iter::repeat_n(0, seq - row.len())));
  }
  let input_ids = Tensor::from_vec(ids, (rows.len(), seq), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(mask, (rows.len(), seq), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let lengths: Vec<usize> = rows.iter().map(Vec::len).collect();
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  model.forward(&input).unwrap().to_vec3::<f32>().unwrap()
}

fn max_shift(a: &[Vec<Vec<f32>>], b: &[Vec<Vec<f32>>]) -> f32 {
  a.iter()
    .flatten()
    .flatten()
    .zip(b.iter().flatten().flatten())
    .map(|(x, y)| (x - y).abs())
    .fold(0.0, f32::max)
}

#[test]
fn output_has_hidden_size_per_token() {
  let (_varmap, model) = tiny_model();
  let out = run(&model, &[vec![3, 5, 7, 9]]);
  assert_eq!((out.len(), out[0].len(), out[0][0].len()), (1, 4, 16));
  assert_eq!(model.hidden_size(), 16);
  assert_eq!(
    model.max_position_embeddings(),
    8192,
    "the trained range and the memory cap, not a table height: ALiBi indexes nothing"
  );
}

#[test]
fn padding_does_not_change_real_tokens() {
  let (_varmap, model) = tiny_model();
  let alone = run(&model, &[vec![3, 9, 5]]);
  let batched = run(&model, &[vec![3, 5, 7, 8, 9], vec![3, 9, 5]]);
  for (pos, token) in alone[0].iter().enumerate() {
    for (a, b) in token.iter().zip(&batched[1][pos]) {
      assert!((a - b).abs() < 1e-5, "position {pos}: {a} vs {b}");
    }
  }
}

/// ALiBi is the only thing here that is not permutation-equivariant, and a reversal leaves `-slope * |i - j|` unchanged,
/// so the test swaps two neighbours instead. Dropping ALiBi costs cosine 0.9423 on `small-en`.
#[test]
fn permuting_two_tokens_is_not_a_permutation_of_the_output() {
  let (_varmap, model) = tiny_model();
  let straight = run(&model, &[vec![3, 5, 7, 9, 11, 13]]);
  let swapped = run(&model, &[vec![3, 7, 5, 9, 11, 13]]);
  let undone: Vec<Vec<f32>> = [0, 2, 1, 3, 4, 5]
    .into_iter()
    .map(|i| swapped[0][i].clone())
    .collect();
  let shift = max_shift(&[straight[0].clone()], &[undone]);
  assert!(
    shift > 1e-3,
    "alibi is the only positional signal, so exchanging two tokens merely exchanged their rows: max shift {shift}"
  );
}

/// Row 0 of the token-type table reaches every token; dropping it costs cosine 0.9333 on `small-en`.
#[test]
fn the_token_type_row_reaches_the_sum() {
  let (mut varmap, model) = tiny_model();
  let rows = [vec![3u32, 5, 7, 9]];
  let before = run(&model, &rows);
  let mut values = vec![0f32; 2 * 16];
  values[3] = 7.0;
  varmap
    .set_one(
      "embeddings.token_type_embeddings.weight",
      Tensor::from_vec(values, (2, 16), &Device::Cpu).unwrap(),
    )
    .unwrap();
  assert!(
    max_shift(&run(&model, &rows), &before) > 1e-3,
    "row 0 of the token type table is read for every token and must reach the sum"
  );
}