candding 0.1.0

//! `pooled_attention` is this head's only tensor code, so this is the one place a real tiny `BertModel` is built.

use std::collections::HashMap;

use candle_core::Device;
use candle_core::Tensor;
use candle_nn::VarBuilder;

use super::*;
use crate::layers::activation::Activation;
use crate::models::bert::BertConfig;

/// Everything but the query and key projections is zeroed, so the expected value needs only those, the LayerNorm and attention.
fn tiny_bert(hidden: usize, heads: usize, vocab: usize) -> BertModel {
  let intermediate = hidden;
  let config = BertConfig {
    vocab_size: vocab,
    hidden_size: hidden,
    num_hidden_layers: 1,
    num_attention_heads: heads,
    intermediate_size: intermediate,
    hidden_act: Activation::Gelu,
    max_position_embeddings: 8,
    type_vocab_size: 1,
    layer_norm_eps: 1e-12,
  };
  let zeros = |rows: usize, cols: usize| vec![0f32; rows * cols];
  let ones = |n: usize| vec![1f32; n];
  let identity = |n: usize| -> Vec<f32> {
    let mut m = vec![0f32; n * n];
    for i in 0..n {
      m[i * n + i] = 1.0;
    }
    m
  };
  // One-hot rows keep every id's embedding independent; the padding id's row is zero.
  let mut word_embeddings = zeros(vocab, hidden);
  for id in 0..vocab.saturating_sub(1) {
    word_embeddings[id * hidden + id] = 1.0;
  }
  let device = Device::Cpu;
  let t = |data: Vec<f32>, shape: (usize, usize)| Tensor::from_vec(data, shape, &device).unwrap();
  let t1 = |data: Vec<f32>| Tensor::from_vec(data, hidden, &device).unwrap();
  let tensors: HashMap<String, Tensor> = HashMap::from([
    (
      "embeddings.word_embeddings.weight".to_string(),
      t(word_embeddings, (vocab, hidden)),
    ),
    (
      "embeddings.position_embeddings.weight".to_string(),
      t(zeros(8, hidden), (8, hidden)),
    ),
    (
      "embeddings.token_type_embeddings.weight".to_string(),
      t(zeros(1, hidden), (1, hidden)),
    ),
    ("embeddings.LayerNorm.weight".to_string(), t1(ones(hidden))),
    (
      "embeddings.LayerNorm.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.attention.self.query.weight".to_string(),
      t(identity(hidden), (hidden, hidden)),
    ),
    (
      "encoder.layer.0.attention.self.query.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.attention.self.key.weight".to_string(),
      t(identity(hidden), (hidden, hidden)),
    ),
    (
      "encoder.layer.0.attention.self.key.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.attention.self.value.weight".to_string(),
      t(zeros(hidden, hidden), (hidden, hidden)),
    ),
    (
      "encoder.layer.0.attention.self.value.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.attention.output.dense.weight".to_string(),
      t(zeros(hidden, hidden), (hidden, hidden)),
    ),
    (
      "encoder.layer.0.attention.output.dense.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.attention.output.LayerNorm.weight".to_string(),
      t1(ones(hidden)),
    ),
    (
      "encoder.layer.0.attention.output.LayerNorm.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.intermediate.dense.weight".to_string(),
      t(zeros(hidden, hidden), (hidden, hidden)),
    ),
    (
      "encoder.layer.0.intermediate.dense.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.output.dense.weight".to_string(),
      t(zeros(hidden, hidden), (hidden, hidden)),
    ),
    (
      "encoder.layer.0.output.dense.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
    (
      "encoder.layer.0.output.LayerNorm.weight".to_string(),
      t1(ones(hidden)),
    ),
    (
      "encoder.layer.0.output.LayerNorm.bias".to_string(),
      t1(zeros(1, hidden)),
    ),
  ]);
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &device);
  BertModel::load(config, vb).unwrap()
}

/// Plain slice arithmetic, independent of `layers::layer_norm::LayerNorm`.
fn layer_norm(x: &[f32], eps: f32) -> Vec<f32> {
  let mean = x.iter().sum::<f32>() / x.len() as f32;
  let var = x.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / x.len() as f32;
  let denom = (var + eps).sqrt();
  x.iter().map(|v| (v - mean) / denom).collect()
}

fn dot(a: &[f32], b: &[f32]) -> f32 {
  a.iter().zip(b).map(|(x, y)| x * y).sum()
}

fn softmax(scores: &[f32]) -> Vec<f32> {
  let max = scores.iter().cloned().fold(f32::MIN, f32::max);
  let exps: Vec<f32> = scores.iter().map(|s| (s - max).exp()).collect();
  let sum: f32 = exps.iter().sum();
  exps.iter().map(|e| e / sum).collect()
}

/// The CLS row is read off the query axis and averaged over both heads; the padded key's 0 comes from `AttentionMask::padding`.
#[test]
fn pooled_attention_reads_the_query_axis_and_averages_heads() {
  let hidden = 4;
  let heads = 2;
  let head_dim = hidden / heads;
  let bert = tiny_bert(hidden, heads, 4);
  let bm42 = Bm42Sparse::new(bert, 4, Vec::new());

  let device = Device::Cpu;
  let input_ids = Tensor::from_vec(vec![0u32, 1, 2, 3], (1, 4), &device).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32, 1, 1, 0], (1, 4), &device).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 actual = bm42.pooled_attention(&input).unwrap();
  assert_eq!(actual.len(), 1);
  let row = &actual[0];
  assert_eq!(row.len(), 4);

  // Q = K = identity, so the projected vectors are the LayerNorm output split by head.
  let eps = 1e-12f32;
  let h0 = layer_norm(&[1.0, 0.0, 0.0, 0.0], eps);
  let h1 = layer_norm(&[0.0, 1.0, 0.0, 0.0], eps);
  let h2 = layer_norm(&[0.0, 0.0, 1.0, 0.0], eps);
  let scale = 1.0f32 / (head_dim as f32).sqrt();
  let mut expected = [0f32; 3];
  for head in 0..heads {
    let slice = |h: &[f32]| h[head * head_dim..(head + 1) * head_dim].to_vec();
    let q0 = slice(&h0);
    let keys = [slice(&h0), slice(&h1), slice(&h2)];
    let scores: Vec<f32> = keys.iter().map(|k| scale * dot(&q0, k)).collect();
    let weights = softmax(&scores);
    for (position, weight) in weights.iter().enumerate() {
      expected[position] += weight / heads as f32;
    }
  }

  for (position, &want) in expected.iter().enumerate() {
    let got = row[position];
    assert!(
      (got - want).abs() < 1e-5,
      "position {position}: got {got}, want {want}"
    );
  }
  assert_eq!(row[3], 0.0, "the masked position must be exactly zero");
}