candding 0.1.0

//! Multi-head attention over one fused, biased `qkv_proj`, with rotary `q` and `k` and a biased output.

use candle_core::Module;
use candle_core::Tensor;
use candle_nn::Linear;
use candle_nn::VarBuilder;
use candle_nn::linear;

use super::config::GteNewConfig;
use crate::error::Result;
use crate::layers::attention::scaled_dot_product;
use crate::layers::attention::split_fused_qkv;
use crate::layers::mask::AttentionMask;
use crate::layers::rope::RotaryEmbedding;

pub struct GteNewAttention {
  q: Linear,
  k: Linear,
  v: Linear,
  out: Linear,
  heads: usize,
  head_dim: usize,
  scale: f64,
}

impl GteNewAttention {
  /// `vb` is `encoder.layer.{i}`; the fused bias is narrowed by hand along the rows `split_fused_qkv` uses.
  pub fn load(config: &GteNewConfig, vb: VarBuilder) -> Result<Self> {
    let hidden = config.hidden_size;
    let qkv = linear(hidden, 3 * hidden, vb.pp("attention.qkv_proj"))?;
    let (q_weight, k_weight, v_weight) = split_fused_qkv(qkv.weight(), hidden)?;
    let bias = qkv
      .bias()
      .expect("candle_nn::linear always reads a bias tensor");
    let q_bias = bias.narrow(0, 0, hidden)?.contiguous()?;
    let k_bias = bias.narrow(0, hidden, hidden)?.contiguous()?;
    let v_bias = bias.narrow(0, 2 * hidden, hidden)?.contiguous()?;
    Ok(Self {
      q: Linear::new(q_weight, Some(q_bias)),
      k: Linear::new(k_weight, Some(k_bias)),
      v: Linear::new(v_weight, Some(v_bias)),
      out: linear(hidden, hidden, vb.pp("attention.o_proj"))?,
      heads: config.num_attention_heads,
      head_dim: config.head_dim(),
      scale: 1.0 / (config.head_dim() as f64).sqrt(),
    })
  }

  pub fn forward(
    &self,
    hidden: &Tensor,
    mask: &AttentionMask,
    rope: &RotaryEmbedding,
  ) -> Result<Tensor> {
    let (batch, seq, _) = hidden.dims3()?;
    let split = |xs: Tensor| -> Result<Tensor> {
      Ok(
        xs.reshape((batch, seq, self.heads, self.head_dim))?
          .transpose(1, 2)?
          .contiguous()?,
      )
    };
    let q = split(self.q.forward(hidden)?)?;
    let k = split(self.k.forward(hidden)?)?;
    let v = split(self.v.forward(hidden)?)?;
    let (q, k) = rope.apply(&q, &k)?;
    let context = scaled_dot_product(&q, &k, &v, mask, self.scale)?
      .transpose(1, 2)?
      .reshape((batch, seq, self.heads * self.head_dim))?;
    Ok(self.out.forward(&context)?)
  }
}

#[cfg(test)]
mod tests {
  use std::collections::HashMap;

  use candle_core::DType;
  use candle_core::Device;
  use candle_nn::VarBuilder;

  use super::*;

  fn config_at_hidden(hidden_size: usize) -> GteNewConfig {
    GteNewConfig {
      hidden_size,
      num_hidden_layers: 1,
      num_attention_heads: 1,
      intermediate_size: hidden_size,
      vocab_size: 10,
      type_vocab_size: 0,
      hidden_act: "gelu".to_string(),
      layer_norm_type: "layer_norm".to_string(),
      layer_norm_eps: 1e-12,
      position_embedding_type: "rope".to_string(),
      rope_theta: 10_000.0,
      rope_scaling: None,
      max_position_embeddings: 32,
      pack_qkv: true,
      unpad_inputs: false,
      use_memory_efficient_attention: false,
    }
  }

  fn full_tensors(hidden: usize) -> HashMap<String, Tensor> {
    let mut tensors = HashMap::new();
    tensors.insert(
      "attention.qkv_proj.weight".to_string(),
      Tensor::zeros((3 * hidden, hidden), DType::F32, &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "attention.qkv_proj.bias".to_string(),
      Tensor::zeros(3 * hidden, DType::F32, &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "attention.o_proj.weight".to_string(),
      Tensor::zeros((hidden, hidden), DType::F32, &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "attention.o_proj.bias".to_string(),
      Tensor::zeros(hidden, DType::F32, &Device::Cpu).unwrap(),
    );
    tensors
  }

  /// This family's `qkv_proj` carries a bias, unlike ModernBERT's and NomicBERT's, so a missing one must fail the load.
  #[test]
  fn load_fails_without_the_qkv_bias() {
    let mut tensors = full_tensors(2);
    tensors.remove("attention.qkv_proj.bias");
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    assert!(GteNewAttention::load(&config_at_hidden(2), vb).is_err());
  }

  /// A fused bias of `[1, 1, 2, 2, 3, 3]` at hidden 2 gives query `[1, 1]`, key `[2, 2]` and value `[3, 3]`.
  #[test]
  fn the_fused_bias_splits_into_the_same_thirds_as_the_weight() {
    let mut tensors = full_tensors(2);
    tensors.insert(
      "attention.qkv_proj.bias".to_string(),
      Tensor::new(&[1f32, 1.0, 2.0, 2.0, 3.0, 3.0], &Device::Cpu).unwrap(),
    );
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    let attention = GteNewAttention::load(&config_at_hidden(2), vb).unwrap();
    assert_eq!(
      attention.q.bias().unwrap().to_vec1::<f32>().unwrap(),
      vec![1.0, 1.0]
    );
    assert_eq!(
      attention.k.bias().unwrap().to_vec1::<f32>().unwrap(),
      vec![2.0, 2.0]
    );
    assert_eq!(
      attention.v.bias().unwrap().to_vec1::<f32>().unwrap(),
      vec![3.0, 3.0]
    );
  }
}