candding 0.1.0

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

use super::config::JinaBertConfig;
use super::variant::Variant;
use crate::error::Result;
use crate::layers::attention::scaled_dot_product;
use crate::layers::layer_norm::LayerNorm;
use crate::layers::mask::AttentionMask;

/// Requests a QK norm only where the header has one; a made-up all-ones norm would renormalize every projection.
fn optional_norm(
  present: bool,
  name: &str,
  hidden: usize,
  eps: f64,
  vb: &VarBuilder,
) -> Result<Option<LayerNorm>> {
  if !present {
    return Ok(None);
  }
  Ok(Some(LayerNorm::load(vb.pp(name), hidden, eps)?))
}

fn project(linear: &Linear, norm: Option<&LayerNorm>, hidden: &Tensor) -> Result<Tensor> {
  let projected = linear.forward(hidden)?;
  match norm {
    Some(norm) => Ok(norm.forward(&projected)?),
    None => Ok(projected),
  }
}

/// BERT's attention tensor names with ALiBi in every score and, for two checkpoints, QK norms.
pub struct JinaBertAttention {
  query: Linear,
  key: Linear,
  value: Linear,
  output: Linear,
  /// `base-es` and `base-code` only, and `None` for the base shape rather than an identity norm.
  norm_q: Option<LayerNorm>,
  norm_k: Option<LayerNorm>,
  layer_norm: LayerNorm,
  num_heads: usize,
  head_dim: usize,
}

impl JinaBertAttention {
  /// `vb` is `encoder.layer.{i}.attention`; `variant` comes from the safetensors header, not `config.json`.
  pub fn load(config: &JinaBertConfig, variant: Variant, vb: VarBuilder) -> Result<Self> {
    let hidden = config.hidden_size;
    let eps = config.layer_norm_eps;
    Ok(Self {
      query: linear(hidden, hidden, vb.pp("self.query"))?,
      key: linear(hidden, hidden, vb.pp("self.key"))?,
      value: linear(hidden, hidden, vb.pp("self.value"))?,
      output: linear(hidden, hidden, vb.pp("output.dense"))?,
      norm_q: optional_norm(variant.qk_norm, "self.layer_norm_q", hidden, eps, &vb)?,
      norm_k: optional_norm(variant.qk_norm, "self.layer_norm_k", hidden, eps, &vb)?,
      layer_norm: LayerNorm::load(vb.pp("output.LayerNorm"), hidden, eps)?,
      num_heads: config.num_attention_heads,
      head_dim: config.head_dim(),
    })
  }

  /// The norm runs over the whole projection, before the heads are split.
  pub(super) fn projected_q(&self, hidden: &Tensor) -> Result<Tensor> {
    project(&self.query, self.norm_q.as_ref(), hidden)
  }

  pub(super) fn projected_k(&self, hidden: &Tensor) -> Result<Tensor> {
    project(&self.key, self.norm_k.as_ref(), hidden)
  }

  /// The value is never normalized. `mask` carries ALiBi and is added after the scale, as the reference does;
  /// adding it before would shrink the bias by `sqrt(head_dim)`.
  pub fn forward(&self, hidden: &Tensor, mask: &AttentionMask) -> Result<Tensor> {
    let (batch, seq, _) = hidden.dims3()?;
    let split = |xs: Tensor| -> Result<Tensor> {
      Ok(
        xs.reshape((batch, seq, self.num_heads, self.head_dim))?
          .transpose(1, 2)?
          .contiguous()?,
      )
    };
    let q = split(self.projected_q(hidden)?)?;
    let k = split(self.projected_k(hidden)?)?;
    let v = split(self.value.forward(hidden)?)?;
    let scale = 1.0 / (self.head_dim as f64).sqrt();
    let context = scaled_dot_product(&q, &k, &v, mask, scale)?
      .transpose(1, 2)?
      .reshape((batch, seq, self.num_heads * self.head_dim))?;
    let projected = self.output.forward(&context)?;
    Ok(self.layer_norm.forward(&(projected + hidden)?)?)
  }
}

#[cfg(test)]
mod tests {
  use candle_core::DType;
  use candle_core::Device;
  use candle_nn::VarMap;

  use super::*;

  /// `layer_norm_q` runs over the full hidden width before the head split: the ramp `0..16` gives a first half of
  /// mean -0.8677, where a per-head norm leaves every head at mean 0.
  #[test]
  fn the_qk_norm_runs_over_the_full_hidden_width() {
    let config = JinaBertConfig {
      vocab_size: 32,
      hidden_size: 16,
      num_hidden_layers: 1,
      num_attention_heads: 2,
      intermediate_size: 32,
      type_vocab_size: 2,
      layer_norm_eps: 1e-12,
      max_position_embeddings: 8192,
      position_embedding_type: "alibi".to_string(),
      feed_forward_type: "geglu".to_string(),
    };
    let mut varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
    let attention = JinaBertAttention::load(
      &config,
      Variant {
        qk_norm: true,
        split_norms: false,
      },
      vb,
    )
    .unwrap();
    varmap
      .set_one(
        "self.query.weight",
        Tensor::eye(16, DType::F32, &Device::Cpu).unwrap(),
      )
      .unwrap();
    varmap
      .set_one(
        "self.query.bias",
        Tensor::zeros(16, DType::F32, &Device::Cpu).unwrap(),
      )
      .unwrap();
    varmap
      .set_one(
        "self.layer_norm_q.weight",
        Tensor::ones(16, DType::F32, &Device::Cpu).unwrap(),
      )
      .unwrap();
    varmap
      .set_one(
        "self.layer_norm_q.bias",
        Tensor::zeros(16, DType::F32, &Device::Cpu).unwrap(),
      )
      .unwrap();
    let ramp: Vec<f32> = (0..16).map(|i| i as f32).collect();
    let hidden = Tensor::from_vec(ramp, (1, 1, 16), &Device::Cpu).unwrap();
    let row = attention
      .projected_q(&hidden)
      .unwrap()
      .flatten_all()
      .unwrap()
      .to_vec1::<f32>()
      .unwrap();
    let first_half: f32 = row[..8].iter().sum::<f32>() / 8.0;
    assert!(
      (first_half + 0.8677).abs() < 1e-3,
      "the first half's mean is {first_half}; a per-head normalization would make it exactly 0"
    );
  }
}