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;

mod conventions;
mod load;

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: 2,
    layer_norm_eps: 1e-12,
  }
}

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

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

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

/// Rows are checked by content, since a count cannot tell which end survived.
#[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();
  // CLS (kept), the marker, the real content and SEP.
  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<_>>()
      );
    }
  }
}

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

/// The punctuation row sits between real rows, so no prefix or suffix cut can produce the expected result.
#[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,
    "one of five positions is punctuation and must be dropped"
  );
  let root2 = 2f32.sqrt();
  let expected_without_the_middle_row = [
    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_without_the_middle_row.iter()) {
    for (g, w) in got.iter().zip(want) {
      assert!(
        (g - w).abs() < 1e-6,
        "the surviving rows must skip exactly the punctuation position, not a contiguous end: {:?}",
        rows[0].rows().collect::<Vec<_>>()
      );
    }
  }
}

#[test]
fn forward_query_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, including the ones that would read as padding elsewhere"
  );
}

#[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);
  // No row is all-zero: with no bias to shift it, a zero row cannot become unit length.
  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);
  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);
}