candding 0.1.0

//! A zero-layer model's hidden state is `norm(E[t])`, so every expected score is `norm(E[t]) · (W[yes] - W[no])`.

use std::collections::BTreeMap;
use std::collections::HashMap;

use candle_nn::VarMap;

use super::*;

const VOCAB: usize = 16;
const HIDDEN: usize = 4;
const YES: u32 = 9;
const NO: u32 = 2;

fn config(layers: usize, tied: bool) -> String {
  format!(
    r#"{{"vocab_size":{VOCAB},"hidden_size":{HIDDEN},"num_hidden_layers":{layers},"num_attention_heads":2,
    "num_key_value_heads":1,"head_dim":2,"intermediate_size":8,"max_position_embeddings":32,
    "rms_norm_eps":1e-6,"tie_word_embeddings":{tied}}}"#
  )
}

/// Row `i` of a `(VOCAB, HIDDEN)` matrix whose rows all differ: `seed` separates the embedding from `lm_head`.
fn matrix(seed: f32) -> Vec<f32> {
  (0..VOCAB * HIDDEN)
    .map(|i| ((i as f32 * 0.37 + seed).sin() * 1.5).round() / 2.0 + 0.1 * (i % 3) as f32)
    .collect()
}

fn tensors(lm_head: bool) -> HashMap<String, Tensor> {
  let mut tensors = HashMap::new();
  let embed = Tensor::from_vec(matrix(0.0), (VOCAB, HIDDEN), &Device::Cpu).unwrap();
  tensors.insert("model.embed_tokens.weight".to_string(), embed);
  let norm = Tensor::ones(HIDDEN, DType::F32, &Device::Cpu).unwrap();
  tensors.insert("model.norm.weight".to_string(), norm);
  if lm_head {
    let head = Tensor::from_vec(matrix(1.3), (VOCAB, HIDDEN), &Device::Cpu).unwrap();
    tensors.insert("lm_head.weight".to_string(), head);
  }
  tensors
}

fn load(
  tied: bool,
  lm_head: bool,
  ids: [u32; 2],
  activation: ScoreActivation,
) -> Result<LogitScoreRerank> {
  let declared = if tied {
    Unembedding::Tied
  } else {
    Unembedding::Untied
  };
  load_declared(declared, tied, lm_head, ids, activation)
}

fn load_declared(
  declared: Unembedding,
  tied: bool,
  lm_head: bool,
  ids: [u32; 2],
  activation: ScoreActivation,
) -> Result<LogitScoreRerank> {
  let vb = VarBuilder::from_tensors(tensors(lm_head), DType::F32, &Device::Cpu);
  LogitScoreRerank::load("qwen3", &config(0, tied), vb, ids, declared, activation)
}

fn expected(token: u32, unembedding: &[f32]) -> f64 {
  let row = |m: &[f32], t: u32| -> Vec<f64> {
    let start = t as usize * HIDDEN;
    m[start..start + HIDDEN]
      .iter()
      .map(|&v| f64::from(v))
      .collect()
  };
  let e = row(&matrix(0.0), token);
  let rms = (e.iter().map(|v| v * v).sum::<f64>() / HIDDEN as f64 + 1e-6).sqrt();
  let (yes, no) = (row(unembedding, YES), row(unembedding, NO));
  e.iter()
    .zip(yes.iter().zip(&no))
    .map(|(h, (y, n))| h / rms * (y - n))
    .sum()
}

/// Two rows of unequal length, right-padded with id 0: the second row's last real token is at position 1.
fn scores(model: &dyn RerankModel, lengths: [usize; 2]) -> Vec<f32> {
  let input_ids = Tensor::new(&[[3u32, 5, 7], [11, 13, 0]], &Device::Cpu).unwrap();
  let attention_mask = Tensor::new(&[[1u32, 1, 1], [1, 1, 0]], &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: &lengths,
  };
  let scores = model.forward(&input).unwrap();
  scores.iter().map(RerankScore::value).collect()
}

fn assert_close(actual: f32, expected: f64, label: &str) {
  assert!(
    (f64::from(actual) - expected).abs() < 1e-5,
    "{label}: {actual} vs {expected}"
  );
}

mod tie;

#[test]
fn swapping_the_token_ids_negates_the_score() {
  let model = load(true, false, [YES, NO], ScoreActivation::Identity).unwrap();
  let swapped = load(true, false, [NO, YES], ScoreActivation::Identity).unwrap();
  let (a, b) = (scores(&model, [3, 2]), scores(&swapped, [3, 2]));
  for (a, b) in a.iter().zip(&b) {
    assert_eq!(*a, -*b);
    assert!(a.abs() > 1e-3, "{a}");
  }
}

#[test]
fn a_token_id_outside_the_vocabulary_is_refused() {
  let refused = load(true, false, [YES, VOCAB as u32], ScoreActivation::Identity);
  assert!(matches!(
    refused,
    Err(CanddingError::LogitTokenOutOfRange {
      token_id: 16,
      vocab_size: 16
    })
  ));
}

#[test]
fn a_padded_row_is_read_at_its_last_real_position() {
  let model = load(true, false, [YES, NO], ScoreActivation::Identity).unwrap();
  let got = scores(&model, [3, 2]);
  assert_close(got[1], expected(13, &matrix(0.0)), "last real token");
  let last_column = scores(&model, [3, 3]);
  assert_close(
    last_column[1],
    expected(0, &matrix(0.0)),
    "the padding position",
  );
  assert!((got[1] - last_column[1]).abs() > 1e-3);
}

#[test]
fn the_activation_applies_to_the_difference() {
  let identity = load(true, false, [YES, NO], ScoreActivation::Identity).unwrap();
  let sigmoid = load(true, false, [YES, NO], ScoreActivation::Sigmoid).unwrap();
  for (raw, squashed) in scores(&identity, [3, 2])
    .iter()
    .zip(scores(&sigmoid, [3, 2]))
  {
    assert!((squashed - 1.0 / (1.0 + (-raw).exp())).abs() < 1e-6);
  }
}

/// The tensors take values in name order: `VarMap`'s `HashMap` order varies per process and some orders scored the padding within `1e-3`.
#[test]
fn a_padded_row_scores_what_it_scores_alone_through_a_real_layer() {
  let varmap = VarMap::new();
  let config = config(1, true);
  LogitScoreRerank::load(
    "qwen3",
    &config,
    VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu),
    [YES, NO],
    Unembedding::Tied,
    ScoreActivation::Identity,
  )
  .unwrap();
  {
    let data = varmap.data().lock().unwrap();
    let by_name: BTreeMap<_, _> = data.iter().collect();
    for (index, var) in by_name.values().enumerate() {
      let values: Vec<f32> = (0..var.elem_count())
        .map(|i| ((i * 7 + index * 13) as f32 * 0.61).sin() * 0.5)
        .collect();
      var
        .set(&Tensor::from_vec(values, var.shape(), &Device::Cpu).unwrap())
        .unwrap();
    }
  }
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = LogitScoreRerank::load(
    "qwen3",
    &config,
    vb,
    [YES, NO],
    Unembedding::Tied,
    ScoreActivation::Identity,
  )
  .unwrap();
  let batched = scores(&model, [3, 2]);
  let input_ids = Tensor::new(&[[11u32, 13]], &Device::Cpu).unwrap();
  let mask = input_ids.ones_like().unwrap();
  let types = input_ids.zeros_like().unwrap();
  let alone = model
    .forward(&EncoderInput {
      input_ids: &input_ids,
      attention_mask: &mask,
      token_type_ids: &types,
      lengths: &[2],
    })
    .unwrap()[0]
    .value();
  assert!(
    (batched[1] - alone).abs() < 1e-5,
    "{} vs {alone}",
    batched[1]
  );
  assert!((scores(&model, [3, 3])[1] - alone).abs() > 1e-3);
}

/// At F16 the logits are 4000 and 0.5: subtracted in F16 they give 4000, where an upcast first gives 3999.5.
#[test]
fn the_logits_are_subtracted_in_the_model_dtype_before_the_upcast() {
  let embed: Vec<f32> = (0..VOCAB)
    .flat_map(|t| match t as u32 {
      YES => [1000.0, 1000.0, 1000.0, 1000.0],
      NO => [0.125, 0.125, 0.125, 0.125],
      _ => [1.0, 1.0, 1.0, 1.0],
    })
    .collect();
  let mut tensors = HashMap::new();
  let embed = Tensor::from_vec(embed, (VOCAB, HIDDEN), &Device::Cpu).unwrap();
  tensors.insert("model.embed_tokens.weight".to_string(), embed);
  let norm = Tensor::ones(HIDDEN, DType::F32, &Device::Cpu).unwrap();
  tensors.insert("model.norm.weight".to_string(), norm);
  let vb = VarBuilder::from_tensors(tensors, DType::F16, &Device::Cpu);
  let model = LogitScoreRerank::load(
    "qwen3",
    &config(0, true),
    vb,
    [YES, NO],
    Unembedding::Tied,
    ScoreActivation::Identity,
  )
  .unwrap();
  let [_, second] = scores(&model, [3, 2])[..] else {
    unreachable!()
  };
  assert_eq!(second, 4000.0);
}