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::BertConfig;

mod conventions;
mod load;
mod pickle;

fn tiny_config(hidden_size: usize) -> BertConfig {
  BertConfig {
    vocab_size: 64,
    hidden_size,
    num_hidden_layers: 0,
    num_attention_heads: 1,
    intermediate_size: 1,
    hidden_act: Activation::Gelu,
    max_position_embeddings: 32,
    type_vocab_size: 1,
    layer_norm_eps: 1e-12,
  }
}

/// `BAAI/bge-m3`'s own `pad_token_id` (1) and `position_embedding_type` ("absolute").
fn xlm_config(bert: BertConfig) -> XlmRobertaConfig {
  XlmRobertaConfig {
    bert,
    pad_token_id: 1,
    position_embedding_type: "absolute".to_string(),
  }
}

/// Never run: every `project_rows` test builds `hidden` by hand.
fn dummy_encoder(hidden_size: usize) -> XlmRobertaModel {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  XlmRobertaModel::load(xlm_config(tiny_config(hidden_size)), vb).unwrap()
}

fn head(weight: Tensor, bias: Tensor, dim: usize, keeps_cls_row: bool) -> M3Colbert {
  M3Colbert {
    encoder: dummy_encoder(weight.dim(1).unwrap()),
    weight,
    bias: Some(bias),
    dim,
    normalize: true,
    keeps_cls_row,
  }
}

fn identity_head(dim: usize) -> M3Colbert {
  let weight = Tensor::eye(dim, DType::F32, &Device::Cpu).unwrap();
  let bias = Tensor::zeros(dim, DType::F32, &Device::Cpu).unwrap();
  head(weight, bias, dim, false)
}

/// `projection_in: 4` matches `tiny_config(4)`; the rest stay at `BAAI/bge-m3`'s real values.
fn conventions(keeps_cls_row: bool) -> MultiVectorConventions {
  MultiVectorConventions {
    keeps_cls_row,
    projection_in: 4,
    projection_out: 4,
    projection_bias: true,
    normalize: true,
    attends_to_mask_tokens: false,
    position_scheme: None,
  }
}

/// Rows are checked by content, since dropping either end of an unpadded row leaves the same count.
#[test]
fn the_classification_row_is_dropped_not_the_final_row() {
  let model = identity_head(2);
  let hidden = Tensor::new(
    &[[[1f32, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, -1.0]]],
    &Device::Cpu,
  )
  .unwrap();
  let rows = model.project_rows(&hidden, &[4]).unwrap();

  assert_eq!(
    rows[0].row_count(),
    3,
    "4 real tokens minus the dropped classification row is 3"
  );
  let root2 = 2f32.sqrt();
  // A, B and SEP, normalized.
  let expected = [
    vec![0.0f32, 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<_>>()
      );
    }
  }
}

/// `<s></s>`, the `empty` fixture's real length of 2, still yields one row: the end token survives the CLS drop.
#[test]
fn the_minimum_real_length_still_yields_one_row() {
  let model = identity_head(2);
  let hidden = Tensor::new(&[[[1f32, 0.0], [0.0, 1.0]]], &Device::Cpu).unwrap();
  let rows = model.project_rows(&hidden, &[2]).unwrap();
  assert_eq!(rows[0].row_count(), 1);
  let only: Vec<f32> = rows[0].rows().next().unwrap().to_vec();
  assert!(
    (only[0] - 0.0).abs() < 1e-6 && (only[1] - 1.0).abs() < 1e-6,
    "{only:?}"
  );
}

#[test]
fn every_row_is_unit_length_after_projection() {
  let weight = Tensor::new(&[[3f32, 0.0], [0.0, 4.0]], &Device::Cpu).unwrap();
  let bias = Tensor::new(&[1f32, 1.0], &Device::Cpu).unwrap();
  let model = head(weight, bias, 2, false);
  let hidden = Tensor::new(&[[[2f32, 2.0], [5.0, -1.0], [0.0, 0.0]]], &Device::Cpu).unwrap();
  let rows = model.project_rows(&hidden, &[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}");
  }
}

/// Padding holds an implausible `[9, 9]`, so any leak into the output is obvious.
#[test]
fn padding_never_reaches_the_output() {
  let model = identity_head(2);
  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 rows = model.project_rows(&hidden, &[4, 2]).unwrap();

  assert_eq!(rows[0].row_count(), 3);
  assert_eq!(rows[1].row_count(), 1);
  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"
  );
  let only_of_second: Vec<f32> = rows[1].rows().next().unwrap().to_vec();
  assert!(
    (only_of_second[1] - 1.0).abs() < 1e-6,
    "{only_of_second:?} leaked padding"
  );
}