candding 0.1.0

use std::collections::HashSet;

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

use super::*;
use crate::layers::activation::Activation;
use crate::models::xlm_roberta_flash::PositionScheme;

mod conventions;
mod load;
mod mask_attendance;

fn tiny_config(hidden_size: usize) -> XlmRobertaFlashConfig {
  XlmRobertaFlashConfig {
    vocab_size: 64,
    hidden_size,
    num_hidden_layers: 0,
    num_attention_heads: 1,
    intermediate_size: 1,
    type_vocab_size: 1,
    max_position_embeddings: 32,
    layer_norm_eps: 1e-5,
    position_embedding_type: Some("rotary".to_string()),
    pad_token_id: 1,
    rotary_emb_base: 10_000.0,
    hidden_act: Activation::Gelu,
    lora_adaptations: None,
    load_trained_adapters: false,
    matryoshka_dimensions: None,
  }
}

/// Never run: the row-selection tests build `hidden` by hand.
fn dummy_encoder(hidden_size: usize) -> XlmRobertaFlashModel {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  XlmRobertaFlashModel::load(tiny_config(hidden_size), vb, Some(PositionScheme::Rotary)).unwrap()
}

fn head(
  weight: Tensor,
  punctuation_ids: HashSet<u32>,
  dim: usize,
  keeps_cls_row: bool,
  attends_to_mask_tokens: bool,
) -> JinaColbert {
  JinaColbert {
    encoder: dummy_encoder(weight.dim(1).unwrap()),
    weight,
    bias: None,
    punctuation_ids,
    dim,
    normalize: true,
    keeps_cls_row,
    attends_to_mask_tokens,
  }
}

fn identity_head(dim: usize, punctuation_ids: HashSet<u32>) -> JinaColbert {
  let weight = Tensor::eye(dim, DType::F32, &Device::Cpu).unwrap();
  head(weight, punctuation_ids, dim, true, true)
}

/// Rows are checked by content, since dropping either end leaves the same count.
#[test]
fn the_classification_row_is_kept_not_dropped() {
  let model = identity_head(2, HashSet::new());
  let hidden = Tensor::new(
    &[[[1f32, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0]]],
    &Device::Cpu,
  )
  .unwrap();
  let ids = vec![vec![2u32, 16, 4, 3]];
  let rows = model.select_document_rows(&hidden, &ids, &[4]).unwrap();
  assert_eq!(rows[0].row_count(), 4, "all four positions survive");
  let root2 = 2f32.sqrt();
  let expected = [
    vec![1.0f32, 0.0],
    vec![0.0, 1.0],
    vec![1.0 / root2, 1.0 / root2],
    vec![1.0 / root2, -1.0 / root2],
  ];
  for (got, want) in rows[0].rows().zip(expected.iter()) {
    for (g, w) in got.iter().zip(want) {
      assert!(
        (g - w).abs() < 1e-6,
        "{:?} vs {expected:?}",
        rows[0].rows().collect::<Vec<_>>()
      );
    }
  }
}

#[test]
fn padding_never_reaches_the_output() {
  let model = identity_head(2, HashSet::new());
  let hidden = Tensor::new(
    &[
      [[1f32, 0.0], [1.0, 1.0], [2.0, 1.0], [3.0, 1.0], [9.0, 9.0]],
      [[1f32, 0.0], [0.0, 5.0], [9.0, 9.0], [9.0, 9.0], [9.0, 9.0]],
    ],
    &Device::Cpu,
  )
  .unwrap();
  let ids = vec![vec![2u32, 5, 5, 5, 0], vec![2u32, 5, 0, 0, 0]];
  let rows = model.select_document_rows(&hidden, &ids, &[4, 2]).unwrap();
  assert_eq!(rows[0].row_count(), 4);
  assert_eq!(rows[1].row_count(), 2);
  let last_of_first: Vec<f32> = rows[0].rows().last().unwrap().to_vec();
  assert!(
    (last_of_first[0] - 3.0 / 10f32.sqrt()).abs() < 1e-5,
    "{last_of_first:?} leaked padding"
  );
}

#[test]
fn punctuation_is_dropped_by_a_gather_not_a_truncation() {
  let punctuation_token_id = 99u32;
  let mut punctuation_ids = HashSet::new();
  punctuation_ids.insert(punctuation_token_id);
  let model = identity_head(2, punctuation_ids);
  let hidden = Tensor::new(
    &[[[1f32, 0.0], [0.0, 1.0], [1.0, 1.0], [5.0, 5.0], [1.0, -1.0]]],
    &Device::Cpu,
  )
  .unwrap();
  let ids = vec![vec![2u32, 16, 4, punctuation_token_id, 3]];
  let rows = model.select_document_rows(&hidden, &ids, &[5]).unwrap();
  assert_eq!(
    rows[0].row_count(),
    4,
    "the punctuation position must be dropped"
  );
  let root2 = 2f32.sqrt();
  let expected = [
    vec![1.0f32, 0.0],
    vec![0.0, 1.0],
    vec![1.0 / root2, 1.0 / root2],
    vec![1.0 / root2, -1.0 / root2],
  ];
  for (got, want) in rows[0].rows().zip(expected.iter()) {
    for (g, w) in got.iter().zip(want) {
      assert!(
        (g - w).abs() < 1e-6,
        "{:?} vs {expected:?}",
        rows[0].rows().collect::<Vec<_>>()
      );
    }
  }
}

#[test]
fn select_query_rows_keeps_every_row_unconditionally() {
  let model = identity_head(2, {
    let mut set = HashSet::new();
    set.insert(99u32);
    set
  });
  let hidden = Tensor::new(
    &[[[1f32, 0.0], [0.0, 1.0], [9.0, 9.0], [9.0, 9.0]]],
    &Device::Cpu,
  )
  .unwrap();
  let rows = model.select_query_rows(&hidden).unwrap();
  assert_eq!(rows[0].row_count(), 4, "every position survives");
}

#[test]
fn every_row_is_unit_length_after_projection() {
  let weight = Tensor::new(&[[3f32, 0.0], [0.0, 4.0]], &Device::Cpu).unwrap();
  let model = head(weight, HashSet::new(), 2, true, true);
  let hidden = Tensor::new(&[[[2f32, 2.0], [5.0, -1.0], [1.0, -3.0]]], &Device::Cpu).unwrap();
  let ids = vec![vec![2u32, 5, 3]];
  let rows = model.select_document_rows(&hidden, &ids, &[3]).unwrap();
  for row in rows[0].rows() {
    let norm = row.iter().map(|v| v * v).sum::<f32>().sqrt();
    assert!((norm - 1.0).abs() < 1e-5, "{row:?} has norm {norm}");
  }
}

#[test]
fn dim_is_read_from_the_weights_own_shape() {
  let weight = Tensor::from_vec(
    (0..12).map(|v| v as f32 * 0.1).collect(),
    (3, 4),
    &Device::Cpu,
  )
  .unwrap();
  let model = head(weight, HashSet::new(), 3, true, true);
  assert_eq!(model.dim(), 3);
  let hidden = Tensor::new(&[[[1f32, 0.0, 0.0, 0.0]]], &Device::Cpu).unwrap();
  let rows = model.select_query_rows(&hidden).unwrap();
  assert_eq!(rows[0].dim(), 3);
}