candding 0.1.0

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use candle_nn::VarBuilder;
use candle_nn::VarMap;

use super::*;
use crate::models::Encoder;
use crate::models::EncoderInput;

fn tiny_config() -> NomicBertConfig {
  NomicBertConfig {
    n_embd: 16,
    n_head: 4,
    n_layer: 2,
    n_inner: 32,
    activation_function: "swiglu".to_string(),
    qkv_proj_bias: false,
    mlp_fc1_bias: false,
    mlp_fc2_bias: false,
    prenorm: false,
    causal: false,
    rotary_emb_interleaved: false,
    norm_mlp: false,
    type_vocab_size: 2,
    vocab_size: 32,
    pad_token_id: Some(0),
    layer_norm_epsilon: 1e-12,
    rotary_emb_fraction: 1.0,
    rotary_emb_base: 1000.0,
    rotary_emb_scale_base: None,
    rotary_scaling_factor: Some(2.0),
    max_trained_positions: 8,
    max_position_embeddings: None,
  }
}

fn tiny_model() -> (VarMap, NomicBertModel) {
  let varmap = VarMap::new();
  let config = tiny_config();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = NomicBertModel::load(config, vb).unwrap();
  (varmap, model)
}

/// The padding id is 0, the tokenizer's `[PAD]` for all three checkpoints.
fn run(model: &NomicBertModel, 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);
}

#[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}");
    }
  }
}

/// Rotary is the only positional signal, so without it a neighbour swap would only swap two output rows.
#[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,
    "rotary is the only positional signal, so exchanging two tokens merely exchanged their rows: max shift {shift}"
  );
}

#[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"
  );
}

/// 12 tokens past `max_trained_positions` 8: the dynamic-NTK base is `1000 * (2*12/8 - 1)^2 = 4000`.
#[test]
fn a_batch_longer_than_the_trained_length_still_runs() {
  let (_varmap, model) = tiny_model();
  let out = run(
    &model,
    &[vec![3u32, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25]],
  );
  for row in &out[0] {
    for value in row {
      assert!(value.is_finite(), "{value}");
    }
  }
}

/// Rotary indexes nothing, so a longer request costs attention memory rather than an out-of-bounds read.
#[test]
fn the_length_is_unbounded() {
  let (_varmap, model) = tiny_model();
  assert_eq!(model.max_position_embeddings(), usize::MAX);
}

#[test]
fn an_unsupported_config_is_refused_at_load() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let mut config = tiny_config();
  config.causal = true;
  assert!(matches!(
    NomicBertModel::load(config, vb).unwrap_err(),
    crate::error::CanddingError::UnsupportedArchitecture(_)
  ));
}