candding 0.1.0

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

use super::config::JinaBertConfig;
use crate::error::Result;
use crate::layers::activation::Activation;
use crate::layers::layer_norm::LayerNorm;
use crate::layers::mlp::GatedMlp;

/// Splits the fused `(2I, H)` `mlp.gated_layers.weight` into the two halves `GatedMlp` wants.
/// `swapped` gates the second half (`base-code`) instead of the first, as each checkpoint does.
pub fn gate_and_up(fused: &Tensor, intermediate: usize, swapped: bool) -> Result<(Tensor, Tensor)> {
  let first = fused.narrow(0, 0, intermediate)?.contiguous()?;
  let second = fused.narrow(0, intermediate, intermediate)?.contiguous()?;
  Ok(if swapped {
    (second, first)
  } else {
    (first, second)
  })
}

/// The gated feed-forward; `act` comes from `feed_forward_type`, not `hidden_act`. The base and QK shapes keep
/// the residual and norm inside this block; the split-norm shape has neither.
pub struct JinaBertMlp {
  gated: GatedMlp,
  layer_norm: Option<LayerNorm>,
}

impl JinaBertMlp {
  /// `vb` is `encoder.layer.{i}`; `split_norms` picks both the tensor names and which half is gated.
  pub fn load(config: &JinaBertConfig, split_norms: bool, vb: VarBuilder) -> Result<Self> {
    let intermediate = config.intermediate_size;
    let (fused_name, down_name) = match split_norms {
      true => ("mlp.up_gated_layer.weight", "mlp.down_layer"),
      false => ("mlp.gated_layers.weight", "mlp.wo"),
    };
    let fused = vb.get((2 * intermediate, config.hidden_size), fused_name)?;
    let (gate, up) = gate_and_up(&fused, intermediate, split_norms)?;
    let down = linear(intermediate, config.hidden_size, vb.pp(down_name))?;
    let layer_norm = match split_norms {
      true => None,
      false => Some(LayerNorm::load(
        vb.pp("mlp.layernorm"),
        config.hidden_size,
        config.layer_norm_eps,
      )?),
    };
    Ok(Self {
      gated: GatedMlp::new(
        Linear::new(gate, None),
        Linear::new(up, None),
        down,
        Activation::Gelu,
      ),
      layer_norm,
    })
  }

  pub fn forward(&self, hidden: &Tensor) -> Result<Tensor> {
    let out = self.gated.forward(hidden)?;
    match &self.layer_norm {
      Some(norm) => Ok(norm.forward(&(out + hidden)?)?),
      None => Ok(out),
    }
  }
}

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

  use super::*;

  /// Swapping the gated halves costs cosine 0.6814 on `small-en` and 0.5396 on `base-en`, with no error anywhere.
  #[test]
  fn the_fused_projection_splits_into_the_halves_each_shape_gates() {
    let mut values = vec![1f32; 2 * 3];
    values[3..].fill(2.0);
    let fused = Tensor::from_vec(values, (2, 3), &Device::Cpu).unwrap();
    let (gate, up) = gate_and_up(&fused, 1, false).unwrap();
    assert_eq!(
      gate.to_vec2::<f32>().unwrap(),
      vec![vec![1.0, 1.0, 1.0]],
      "base shape gates the first half"
    );
    assert_eq!(up.to_vec2::<f32>().unwrap(), vec![vec![2.0, 2.0, 2.0]]);
    let (gate, up) = gate_and_up(&fused, 1, true).unwrap();
    assert_eq!(
      gate.to_vec2::<f32>().unwrap(),
      vec![vec![2.0, 2.0, 2.0]],
      "base-code gates the second half"
    );
    assert_eq!(up.to_vec2::<f32>().unwrap(), vec![vec![1.0, 1.0, 1.0]]);
  }

  /// Halves of 1 and 3 give `gelu(1) * 3`, so gating the wrong half gives a different number, not a different shape.
  #[test]
  fn the_block_gates_the_half_it_says_it_does() {
    let hidden = Tensor::new(&[[[1f32]]], &Device::Cpu).unwrap();
    let tensors = std::collections::HashMap::from([
      (
        "mlp.gated_layers.weight".to_string(),
        Tensor::new(&[[1f32], [3.0]], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.wo.weight".to_string(),
        Tensor::new(&[[1f32]], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.wo.bias".to_string(),
        Tensor::new(&[0f32], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.layernorm.weight".to_string(),
        Tensor::new(&[1f32], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.layernorm.bias".to_string(),
        Tensor::new(&[0f32], &Device::Cpu).unwrap(),
      ),
    ]);
    let vb = VarBuilder::from_tensors(tensors, candle_core::DType::F32, &Device::Cpu);
    let config: JinaBertConfig = serde_json::from_str(
      r#"{"vocab_size":8,"hidden_size":1,"num_hidden_layers":1,"num_attention_heads":1,
          "intermediate_size":1,"max_position_embeddings":8,"position_embedding_type":"alibi",
          "feed_forward_type":"geglu","type_vocab_size":2,"layer_norm_eps":1e-12}"#,
    )
    .unwrap();
    let mlp = JinaBertMlp::load(&config, false, vb).unwrap();
    let gated = mlp.gated.forward(&hidden).unwrap();
    let value = gated.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];
    assert!(
      (value - 2.524_034).abs() < 1e-5,
      "gelu(1) * 3 is 2.524034, gating the other half would give gelu(3) * 1 = 2.995950: {value}"
    );
  }

  /// The split shape gates the other half: `gelu(3) * 1` where the base shape gives `gelu(1) * 3`.
  #[test]
  fn the_split_block_gates_the_other_half() {
    let hidden = Tensor::new(&[[[1f32]]], &Device::Cpu).unwrap();
    let tensors = std::collections::HashMap::from([
      (
        "mlp.up_gated_layer.weight".to_string(),
        Tensor::new(&[[1f32], [3.0]], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.down_layer.weight".to_string(),
        Tensor::new(&[[1f32]], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.down_layer.bias".to_string(),
        Tensor::new(&[0f32], &Device::Cpu).unwrap(),
      ),
    ]);
    let vb = VarBuilder::from_tensors(tensors, candle_core::DType::F32, &Device::Cpu);
    let config: JinaBertConfig = serde_json::from_str(
      r#"{"vocab_size":8,"hidden_size":1,"num_hidden_layers":1,"num_attention_heads":1,
          "intermediate_size":1,"max_position_embeddings":8,"position_embedding_type":"alibi",
          "feed_forward_type":"geglu","type_vocab_size":2,"layer_norm_eps":1e-12}"#,
    )
    .unwrap();
    let mlp = JinaBertMlp::load(&config, true, vb).unwrap();
    let gated = mlp.gated.forward(&hidden).unwrap();
    let value = gated.flatten_all().unwrap().to_vec1::<f32>().unwrap()[0];
    assert!(
      (value - 2.995_95).abs() < 1e-5,
      "gelu(3) * 1 is 2.995950, gating the other half would give gelu(1) * 3 = 2.524034: {value}"
    );
  }
}