candding 0.1.0

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

use super::*;
use crate::layers::activation::Activation;
use crate::models::bert::config::BertConfig;
use crate::models::xlm_roberta::embeddings::position_ids;

fn tiny_config() -> XlmRobertaConfig {
  XlmRobertaConfig {
    bert: BertConfig {
      vocab_size: 32,
      hidden_size: 16,
      num_hidden_layers: 2,
      num_attention_heads: 2,
      intermediate_size: 32,
      hidden_act: Activation::Gelu,
      max_position_embeddings: 16,
      type_vocab_size: 1,
      layer_norm_eps: 1e-5,
    },
    pad_token_id: 1,
    position_embedding_type: "absolute".to_string(),
  }
}

/// The `VarMap` comes back with the model so a test can overwrite one tensor in place.
fn tiny_model() -> (VarMap, XlmRobertaModel) {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = XlmRobertaModel::load(tiny_config(), vb).unwrap();
  (varmap, model)
}

fn run(model: &XlmRobertaModel, 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(1, 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 numbered(rows: &[Vec<u32>], pad: u32) -> Vec<Vec<u32>> {
  let width = rows[0].len();
  let flat: Vec<u32> = rows.concat();
  let ids = Tensor::from_vec(flat, (rows.len(), width), &Device::Cpu).unwrap();
  position_ids(&ids, pad).unwrap().to_vec2::<u32>().unwrap()
}

/// Every row's first real token sits at index 2, an `n`-token row's last at `n + 1`, and padding at 1.
#[test]
fn position_ids_start_at_the_pad_offset_and_fall_back_to_pad() {
  let ids = numbered(&[vec![0, 5, 2, 1, 1], vec![0, 5, 7, 9, 2]], 1);
  assert_eq!(ids, vec![vec![2, 3, 4, 1, 1], vec![2, 3, 4, 5, 6]]);
}

/// A typed `<pad>` reads the pad row as padding does, where numbering by row length would read `2..=20`.
#[test]
fn a_pad_id_among_the_real_tokens_reads_the_pad_row_and_moves_the_rest_down() {
  let ids = vec![
    0, 4865, 83, 70, 10323, 111, 9098, 32, 2, 2, 581, 10323, 111, 9098, 83, 1, 134288, 5, 2,
  ];
  let reference = vec![
    2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 17, 18, 19,
  ];
  assert_eq!(numbered(&[ids], 1), vec![reference]);
  let leading = numbered(&[vec![0, 1, 1, 1, 581, 2], vec![0, 5, 1, 7, 2, 1]], 1);
  assert_eq!(
    leading,
    vec![vec![2, 1, 1, 1, 3, 4], vec![2, 3, 1, 4, 5, 1]]
  );
}

#[test]
fn rows_without_a_pad_id_number_as_their_length_did() {
  for pad in [1u32, 0] {
    for lengths in [vec![0, 1, 7], vec![4, 4], vec![1], vec![6, 2, 5, 0]] {
      let seq = *lengths.iter().max().unwrap().max(&1);
      let rows: Vec<Vec<u32>> = lengths
        .iter()
        .enumerate()
        .map(|(row, &length)| {
          (0..seq)
            .map(|i| {
              let id = ((row + 3 * i) % 7) as u32;
              match (i < length, id == pad) {
                (false, _) => pad,
                (true, true) => id + 7,
                (true, false) => id,
              }
            })
            .collect()
        })
        .collect();
      let ids = numbered(&rows, pad);
      for (row, &length) in lengths.iter().enumerate() {
        let by_length: Vec<u32> = (0..seq)
          .map(|i| if i < length { pad + 1 + i as u32 } else { pad })
          .collect();
        assert_eq!(
          ids[row], by_length,
          "pad {pad} lengths {lengths:?} row {row}"
        );
      }
    }
  }
}

/// Row `p` carries its value in column `p`, so a row that reaches the sum survives the layer norm's mean subtraction.
fn write_position_table(varmap: &mut VarMap, loud: &[usize]) {
  let config = tiny_config();
  let (rows, hidden) = (config.bert.max_position_embeddings, config.bert.hidden_size);
  let mut values = vec![0f32; rows * hidden];
  for p in 0..rows {
    values[p * hidden + p % hidden] = if loud.contains(&p) { 9.0 } else { 1.0 };
  }
  varmap
    .set_one(
      "embeddings.position_embeddings.weight",
      Tensor::from_vec(values, (rows, hidden), &Device::Cpu).unwrap(),
    )
    .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)
}

/// Passing `0` for the pad id here breaks none of the `position_ids` tests above.
#[test]
fn the_position_lookup_starts_at_the_configured_offset() {
  let (mut varmap, model) = tiny_model();
  let offset = tiny_config().position_offset() as usize;
  let rows = [vec![0u32, 5, 7, 8, 2]];
  write_position_table(&mut varmap, &[]);
  let base = run(&model, &rows);
  let below: Vec<usize> = (0..offset).collect();
  write_position_table(&mut varmap, &below);
  assert!(
    max_shift(&run(&model, &rows), &base) < 1e-6,
    "rows {below:?} sit below the offset and must never be read"
  );
  write_position_table(&mut varmap, &[offset]);
  assert!(
    max_shift(&run(&model, &rows), &base) > 1e-3,
    "row {offset} is the first row the numbering reaches"
  );
}

#[test]
fn a_real_pad_id_reads_the_tables_pad_row() {
  let (mut varmap, model) = tiny_model();
  let pad = tiny_config().pad_token_id;
  let rows = [vec![0u32, 5, pad as u32, 7, 2]];
  write_position_table(&mut varmap, &[]);
  let base = run(&model, &rows);
  write_position_table(&mut varmap, &[pad]);
  assert!(
    max_shift(&run(&model, &rows), &base) > 1e-3,
    "the real pad id reads row {pad}"
  );
  let past = pad + rows[0].len();
  write_position_table(&mut varmap, &[past]);
  assert!(
    max_shift(&run(&model, &rows), &base) < 1e-6,
    "row {past} is past the row's last number"
  );
}

#[test]
fn output_has_hidden_size_per_token() {
  let (_varmap, model) = tiny_model();
  let out = run(&model, &[vec![0, 5, 7, 2]]);
  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(), 14);
}

/// An offset that drifted by one would index past the table and fail here rather than return a wrong vector.
#[test]
fn a_row_at_the_sequence_limit_still_indexes_the_table() {
  let (_varmap, model) = tiny_model();
  let limit = model.max_position_embeddings();
  let out = run(&model, &[vec![0; limit]]);
  assert_eq!((out.len(), out[0].len()), (1, limit));
}

#[test]
fn padding_does_not_change_real_tokens() {
  let (_varmap, model) = tiny_model();
  let alone = run(&model, &[vec![0, 9, 2]]);
  let batched = run(&model, &[vec![0, 5, 7, 8, 2], vec![0, 9, 2]]);
  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}");
    }
  }
}

/// `type_vocab_size` is 1 for every XLM-RoBERTa checkpoint, so a one-row table must not be a shape error.
#[test]
fn a_one_row_token_type_table_loads_and_runs() {
  let config = tiny_config();
  assert_eq!(config.bert.type_vocab_size, 1);
  let (_varmap, model) = tiny_model();
  let out = run(&model, &[vec![0, 4, 2]]);
  assert_eq!(out[0].len(), 3);
}

#[test]
fn a_relative_position_scheme_is_refused_at_load() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let mut config = tiny_config();
  config.position_embedding_type = "relative_key".to_string();
  assert!(matches!(
    XlmRobertaModel::load(config, vb).unwrap_err(),
    crate::error::CanddingError::UnsupportedArchitecture(_)
  ));
}