candding 0.1.0

//! Pooling of per-token hidden states into one vector per sequence.

use candle_core::DType;
use candle_core::Tensor;
use serde::Deserialize;
use serde::Serialize;

use crate::error::Result;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum Pooling {
  Cls,
  Mean,
  LastToken,
}

/// `hidden` is `(batch, seq, hidden)`, `mask` is `(batch, seq)` in any dtype, `lengths` count real tokens per row.
pub fn pool(pooling: Pooling, hidden: &Tensor, mask: &Tensor, lengths: &[usize]) -> Result<Tensor> {
  match pooling {
    Pooling::Cls => Ok(hidden.narrow(1, 0, 1)?.squeeze(1)?),
    Pooling::Mean => mean(hidden, mask, lengths),
    Pooling::LastToken => last_token(hidden, lengths),
  }
}

/// Mask-weighted mean in F32 over the caller's exact counts: padding is selected away, as `inf * 0` would be NaN,
/// and a half-dtype running sum stalls (F16 at 2048, BF16 at 256), so the mask is never summed.
fn mean(hidden: &Tensor, mask: &Tensor, lengths: &[usize]) -> Result<Tensor> {
  let (batch, _, _) = hidden.dims3()?;
  crate::layers::mask::check_lengths(lengths, batch)?;
  let keep = mask
    .to_dtype(DType::U8)?
    .unsqueeze(2)?
    .broadcast_as(hidden.shape())?;
  let summed = keep.where_cond(hidden, &hidden.zeros_like()?)?.sum(1)?;
  let counts: Vec<f32> = lengths.iter().map(|&len| len as f32).collect();
  let counts = Tensor::from_vec(counts, (batch, 1), hidden.device())?.maximum(1e-9)?;
  Ok(
    summed
      .to_dtype(DType::F32)?
      .broadcast_div(&counts)?
      .to_dtype(hidden.dtype())?,
  )
}

/// Gathers position `length - 1` of every row, which is the last real token under right padding.
fn last_token(hidden: &Tensor, lengths: &[usize]) -> Result<Tensor> {
  let (batch, _, hidden_size) = hidden.dims3()?;
  let index: Vec<u32> = lengths.iter().map(|l| l.saturating_sub(1) as u32).collect();
  let index = Tensor::from_vec(index, (batch, 1, 1), hidden.device())?
    .broadcast_as((batch, 1, hidden_size))?
    .contiguous()?;
  Ok(hidden.gather(&index, 1)?.squeeze(1)?)
}

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

  use super::*;

  fn sample() -> (Tensor, Tensor, Vec<usize>) {
    let hidden = Tensor::new(
      &[
        [[1f32, 2.0], [3.0, 4.0], [5.0, 6.0]],
        [[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]],
      ],
      &Device::Cpu,
    )
    .unwrap();
    let mask = Tensor::new(&[[1u32, 1, 0], [1, 1, 1]], &Device::Cpu).unwrap();
    (hidden, mask, vec![2, 3])
  }

  fn rows(t: Tensor) -> Vec<Vec<f32>> {
    t.to_vec2::<f32>().unwrap()
  }

  #[test]
  fn cls_takes_first_token() {
    let (h, m, l) = sample();
    assert_eq!(
      rows(pool(Pooling::Cls, &h, &m, &l).unwrap()),
      vec![vec![1.0, 2.0], vec![1.0, 1.0]]
    );
  }

  #[test]
  fn mean_ignores_padding() {
    let (h, m, l) = sample();
    assert_eq!(
      rows(pool(Pooling::Mean, &h, &m, &l).unwrap()),
      vec![vec![2.0, 3.0], vec![2.0, 2.0]]
    );
  }

  #[test]
  fn last_token_uses_lengths() {
    let (h, m, l) = sample();
    assert_eq!(
      rows(pool(Pooling::LastToken, &h, &m, &l).unwrap()),
      vec![vec![3.0, 4.0], vec![3.0, 3.0]]
    );
  }

  #[test]
  fn serializes_snake_case() {
    assert_eq!(
      serde_json::to_string(&Pooling::LastToken).unwrap(),
      "\"last_token\""
    );
  }

  #[test]
  fn mean_refuses_a_lengths_slice_that_does_not_match_the_batch() {
    let (h, m, _) = sample();
    for lengths in [&[3usize][..], &[3, 3, 3][..], &[][..]] {
      let err = pool(Pooling::Mean, &h, &m, lengths).unwrap_err();
      assert!(
        matches!(
          err,
          crate::error::CanddingError::LengthsMismatch { rows: 2, .. }
        ),
        "{lengths:?}: {err}"
      );
      assert!(
        err.to_string().contains("one real-token length per row"),
        "{err}"
      );
    }
  }

  #[test]
  fn mean_divides_by_the_true_count_past_the_half_dtype_ceiling() {
    let seq = 3000;
    let mut values = vec![0f32; seq * 2];
    values[0] = seq as f32;
    values[1] = seq as f32;
    let hidden = Tensor::from_vec(values, (1, seq, 2), &Device::Cpu).unwrap();
    let mask = Tensor::ones((1, seq), DType::U32, &Device::Cpu).unwrap();
    let reference = rows(pool(Pooling::Mean, &hidden, &mask, &[seq]).unwrap());
    assert_eq!(reference, vec![vec![1.0, 1.0]]);
    for dtype in [DType::F16, DType::BF16] {
      let half = hidden.to_dtype(dtype).unwrap();
      let pooled = pool(Pooling::Mean, &half, &mask, &[seq])
        .unwrap()
        .to_dtype(DType::F32)
        .unwrap();
      let pooled = rows(pooled);
      assert!(
        (pooled[0][0] - reference[0][0]).abs() < 1e-2,
        "{dtype:?} pooled {pooled:?} against f32 {reference:?}"
      );
    }
  }

  /// A padded position can overflow a half dtype to `inf`, and `inf` times a zero mask is `NaN`.
  #[test]
  fn mean_survives_a_non_finite_padded_position() {
    let hidden = Tensor::from_vec(
      vec![1.0f32, 3.0, f32::INFINITY, f32::INFINITY],
      (1, 2, 2),
      &Device::Cpu,
    )
    .unwrap();
    let mask = Tensor::new(&[[1u32, 0]], &Device::Cpu).unwrap();
    let pooled = pool(Pooling::Mean, &hidden, &mask, &[1])
      .unwrap()
      .to_vec2::<f32>()
      .unwrap();
    assert_eq!(pooled, vec![vec![1.0, 3.0]]);
  }
}