candding 0.1.0

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

use super::*;
use crate::layers::activation::Activation;

/// 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)
  }
}

fn tiny_qwen3_config() -> Qwen3Config {
  Qwen3Config {
    vocab_size: 6,
    hidden_size: 8,
    num_hidden_layers: 1,
    num_attention_heads: 2,
    num_key_value_heads: 1,
    head_dim: 4,
    intermediate_size: 16,
    hidden_act: Activation::Silu,
    max_position_embeddings: 16,
    rms_norm_eps: 1e-6,
    rope_theta: 10000.0,
    attention_bias: false,
    use_sliding_window: false,
    is_causal: false,
    tie_word_embeddings: false,
  }
}

/// Never invoked: the tests that use it call `pool_logits` directly.
fn dummy_qwen3() -> Qwen3Model {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  Qwen3Model::load(tiny_qwen3_config(), vb).unwrap()
}

fn head(embed_tokens: Tensor, chunk: usize) -> SpladeTiedSparse {
  let vocab_size = embed_tokens.dims()[0];
  SpladeTiedSparse {
    qwen3: dummy_qwen3(),
    embed_tokens,
    vocab_size,
    chunk,
  }
}

fn tiny_model() -> SpladeTiedSparse {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  SpladeTiedSparse::load(tiny_qwen3_config(), vb).unwrap()
}

/// `FreshEachTime` makes a second read of `embed_tokens.weight` observable as a distinct `TensorId`.
#[test]
fn load_reads_the_decoder_weight_from_the_tied_embedding() {
  let vb = VarBuilder::from_backend(
    Box::new(FreshEachTime(VarMap::new())),
    DType::F32,
    Device::Cpu,
  );
  let model = SpladeTiedSparse::load(tiny_qwen3_config(), vb).unwrap();

  assert_eq!(model.embed_tokens.id(), model.qwen3.embed_tokens().id());
  assert_eq!(model.vocab_size(), 6);
  assert_eq!(model.max_position_embeddings(), 16);
}

#[test]
fn load_accepts_the_model_prefixed_tensor_names() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu).pp("model");
  let model = SpladeTiedSparse::load(tiny_qwen3_config(), vb).unwrap();
  assert_eq!(model.vocab_size(), 6);
}

fn encoder_input(ids: &[u32], device: &Device) -> (Tensor, Tensor, Tensor) {
  let input_ids = Tensor::from_vec(ids.to_vec(), (1, ids.len()), device).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32; ids.len()], (1, ids.len()), device).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  (input_ids, attention_mask, token_type_ids)
}

#[test]
fn forward_produces_one_row_per_input_with_ascending_positive_output() {
  let model = tiny_model();
  let device = Device::Cpu;
  let (input_ids, attention_mask, token_type_ids) = encoder_input(&[1, 2, 3, 4], &device);
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &[4],
  };
  let rows = model.forward(&input).unwrap();
  assert_eq!(rows.len(), 1);
  assert!(rows[0].indices.windows(2).all(|w| w[0] < w[1]));
  assert!(rows[0].values.iter().all(|&v| v > 0.0));
}

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:?}");
    }
  }
}

#[test]
fn the_chunked_maximum_equals_the_whole() {
  let embed_tokens = 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 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> {
    head(embed_tokens.clone(), chunk)
      .pool_logits(&hidden, &mask)
      .unwrap()
  };

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

#[test]
fn padded_positions_never_contribute() {
  let embed_tokens = Tensor::new(&[[1f32, 0.0, 0.0]], &Device::Cpu).unwrap();
  let model = head(embed_tokens, CHUNK);

  let alone_hidden = Tensor::new(&[[[0.3f32, 0.3, 0.3], [0.31, 0.29, 0.3]]], &Device::Cpu).unwrap();
  let alone_mask = Tensor::new(&[[1u32, 1]], &Device::Cpu).unwrap();
  let padded_hidden = Tensor::new(
    &[[[0.3f32, 0.3, 0.3], [0.31, 0.29, 0.3], [50.0, -50.0, -50.0]]],
    &Device::Cpu,
  )
  .unwrap();
  let padded_mask = Tensor::new(&[[1u32, 1, 0]], &Device::Cpu).unwrap();

  let alone = model.pool_logits(&alone_hidden, &alone_mask).unwrap();
  let padded = model.pool_logits(&padded_hidden, &padded_mask).unwrap();
  assert_rows_close(&alone, &padded);
}

mod empty;