candding 0.1.0

use std::collections::HashMap;

use candle_core::DType;
use candle_core::Device;
use candle_core::Shape;
use candle_nn::VarMap;
use candle_nn::var_builder::SimpleBackend;

use super::*;

/// Returns a fresh tensor on every `get`, as the safetensors loader does; `VarMap::get` would return the same one.
struct FreshEachTime(VarMap);

impl SimpleBackend for FreshEachTime {
  fn get(
    &self,
    s: Shape,
    name: &str,
    h: candle_nn::Init,
    dtype: DType,
    dev: &Device,
  ) -> candle_core::Result<Tensor> {
    self.0.get(s, name, h, dtype, dev)?.affine(1.0, 0.0)
  }

  fn get_unchecked(&self, name: &str, _dtype: DType, _dev: &Device) -> candle_core::Result<Tensor> {
    candle_core::bail!("FreshEachTime::get_unchecked is unused by this test: {name}")
  }

  fn contains_tensor(&self, name: &str) -> bool {
    self.0.data().lock().unwrap().contains_key(name)
  }
}

/// Never invoked: the tests call `pool_logits` directly with hand-set tensors.
fn dummy_bert() -> BertModel {
  let config = BertConfig {
    vocab_size: 1,
    hidden_size: 1,
    num_hidden_layers: 0,
    num_attention_heads: 1,
    intermediate_size: 1,
    hidden_act: Activation::Gelu,
    max_position_embeddings: 1,
    type_vocab_size: 1,
    layer_norm_eps: 1e-12,
  };
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  BertModel::load(config, vb).unwrap()
}

fn norm(weight: Tensor, bias: Tensor, size: usize, eps: f64) -> LayerNorm {
  let tensors = HashMap::from([("weight".to_string(), weight), ("bias".to_string(), bias)]);
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  LayerNorm::load(vb, size, eps).unwrap()
}

fn head(
  transform: Linear,
  transform_norm: LayerNorm,
  decoder_weight: Tensor,
  decoder_bias: Tensor,
  chunk: usize,
) -> SpladeSparse {
  let vocab_size = decoder_weight.dims()[0];
  SpladeSparse {
    bert: dummy_bert(),
    transform,
    transform_norm,
    decoder_weight,
    decoder_bias,
    vocab_size,
    chunk,
  }
}

fn assert_rows_close(a: &[SparseEmbedding], b: &[SparseEmbedding]) {
  assert_eq!(a.len(), b.len(), "row count: {a:?} vs {b:?}");
  for (row_a, row_b) in a.iter().zip(b) {
    assert_eq!(row_a.indices, row_b.indices, "{row_a:?} vs {row_b:?}");
    for (va, vb) in row_a.values.iter().zip(&row_b.values) {
      assert!((va - vb).abs() < 1e-5, "{row_a:?} vs {row_b:?}");
    }
  }
}

/// A logit of 3.0 gives `log1p(3) = 1.3862944`; `relu` alone gives 3.0 and `log` alone 1.0986123.
#[test]
fn every_logit_goes_through_relu_then_log1p() {
  let transform = Linear::new(
    Tensor::new(&[[0f32]], &Device::Cpu).unwrap(),
    Some(Tensor::new(&[0f32], &Device::Cpu).unwrap()),
  );
  // LayerNorm over one feature returns its bias, fixing the pre-decode value at 3.0.
  let transform_norm = norm(
    Tensor::new(&[1f32], &Device::Cpu).unwrap(),
    Tensor::new(&[3f32], &Device::Cpu).unwrap(),
    1,
    1e-12,
  );
  let decoder_weight = Tensor::new(&[[1f32]], &Device::Cpu).unwrap();
  let decoder_bias = Tensor::new(&[0f32], &Device::Cpu).unwrap();
  let model = head(
    transform,
    transform_norm,
    decoder_weight,
    decoder_bias,
    CHUNK,
  );

  let hidden = Tensor::new(&[[[0.5f32]]], &Device::Cpu).unwrap();
  let mask = Tensor::new(&[[1u32]], &Device::Cpu).unwrap();
  let rows = model.pool_logits(&hidden, &mask).unwrap();

  assert_eq!(rows[0].indices, vec![0]);
  let w = rows[0].values[0];
  assert!((w - 1.386_294_4).abs() < 1e-6, "{w}");
}

/// The max is associative, so the head keeps a `(batch, vocab)` running max: 3.9 MB against 1.86 GiB at the default batch.
#[test]
fn the_chunked_maximum_equals_the_whole() {
  let transform_weight = Tensor::new(
    &[[0.3f32, -0.2, 0.1], [0.1, 0.4, -0.3], [-0.2, 0.2, 0.5]],
    &Device::Cpu,
  )
  .unwrap();
  let transform_bias = Tensor::new(&[0.05f32, -0.1, 0.2], &Device::Cpu).unwrap();
  let norm_weight = Tensor::new(&[1.1f32, 0.9, 1.0], &Device::Cpu).unwrap();
  let norm_bias = Tensor::new(&[0.1f32, -0.1, 0.0], &Device::Cpu).unwrap();
  let decoder_weight = Tensor::new(
    &[
      [0.5f32, -0.4, 0.3],
      [0.2, 0.6, -0.5],
      [-0.3, 0.1, 0.4],
      [0.7, -0.2, 0.1],
    ],
    &Device::Cpu,
  )
  .unwrap();
  let decoder_bias = Tensor::new(&[0.3f32, 0.1, -0.05, 0.02], &Device::Cpu).unwrap();
  let hidden = Tensor::new(
    &[[
      [0.2f32, 0.5, -0.3],
      [1.0, -0.4, 0.6],
      [-0.5, 0.3, 0.9],
      [0.4, 0.4, -0.2],
      [-0.1, -0.6, 0.7],
    ]],
    &Device::Cpu,
  )
  .unwrap();
  let mask = Tensor::new(&[[1u32, 1, 1, 1, 1]], &Device::Cpu).unwrap();

  let run = |chunk: usize| -> Vec<SparseEmbedding> {
    let model = head(
      Linear::new(transform_weight.clone(), Some(transform_bias.clone())),
      norm(norm_weight.clone(), norm_bias.clone(), 3, 1e-12),
      decoder_weight.clone(),
      decoder_bias.clone(),
      chunk,
    );
    model.pool_logits(&hidden, &mask).unwrap()
  };

  let whole = run(5);
  assert_rows_close(&whole, &run(3));
  assert_rows_close(&whole, &run(1));
}

/// `FreshEachTime` makes a second read of the embedding table observable as a distinct `TensorId`.
#[test]
fn load_reads_the_decoder_weight_from_the_tied_embedding() {
  let config = BertConfig {
    vocab_size: 5,
    hidden_size: 4,
    num_hidden_layers: 1,
    num_attention_heads: 2,
    intermediate_size: 8,
    hidden_act: Activation::Gelu,
    max_position_embeddings: 8,
    type_vocab_size: 2,
    layer_norm_eps: 1e-12,
  };
  let vb = VarBuilder::from_backend(
    Box::new(FreshEachTime(VarMap::new())),
    DType::F32,
    Device::Cpu,
  );
  let model = SpladeSparse::load(config, vb).unwrap();

  assert_eq!(model.decoder_weight.id(), model.bert.embed_tokens().id());

  let input_ids = Tensor::from_vec(vec![0u32, 1, 2], (1, 3), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32, 1, 1], (1, 3), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &[3],
  };
  let rows = model.forward(&input).unwrap();
  assert_eq!(rows.len(), 1);
  assert_eq!(model.vocab_size(), 5);
  assert_eq!(model.max_position_embeddings(), 8);
}

mod properties;