candding 0.1.0

//! Multi-vector output: one row per surviving token, and the late-interaction score between two such matrices.

/// A matrix of per-token rows in one row-major buffer; `dim` is the row width, kept even when there are no rows.
#[derive(Clone, Debug, PartialEq)]
pub struct MultiVectorEmbedding {
  data: Vec<f32>,
  dim: usize,
}

impl MultiVectorEmbedding {
  /// Builds an embedding from rows of exactly `dim` components each; panics if a row has another length.
  pub fn from_rows(dim: usize, rows: impl IntoIterator<Item = Vec<f32>>) -> Self {
    let mut data = Vec::new();
    for (index, row) in rows.into_iter().enumerate() {
      assert_eq!(
        row.len(),
        dim,
        "multi-vector row {index} has {} components, expected {dim}",
        row.len()
      );
      data.extend(row);
    }
    Self { data, dim }
  }

  /// The width of every row, even when there are no rows.
  pub fn dim(&self) -> usize {
    self.dim
  }

  /// The number of rows.
  pub fn row_count(&self) -> usize {
    self.data.len().checked_div(self.dim).unwrap_or(0)
  }

  /// Whether the embedding has no rows.
  pub fn is_empty(&self) -> bool {
    self.row_count() == 0
  }

  /// The rows, in order.
  pub fn rows(&self) -> impl Iterator<Item = &[f32]> {
    self.data.chunks(self.dim.max(1))
  }
}

/// An empty document's score whatever the query, so it sorts behind every real one: sentence-transformers'
/// fixed `_EMPTY_DOCUMENT_SCORE`, not a sum of `DtypeMinimum` fills; checked before an empty query's `0.0`.
const EMPTY_DOCUMENT_SCORE: f32 = -1e9;

/// The value a masked document row takes before the per-query-row maximum, so padding never wins it.
/// Per-checkpoint data, because the references differ.
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MaskedRowFill {
  /// The dtype's minimum, as sentence-transformers fills it: the three ColBERT-style checkpoints.
  DtypeMinimum,
  /// Zero, as FlagEmbedding masks it: `BAAI/bge-m3`.
  Zero,
}

impl MaskedRowFill {
  fn value(self) -> f32 {
    match self {
      MaskedRowFill::DtypeMinimum => f32::MIN,
      MaskedRowFill::Zero => 0.0,
    }
  }
}

/// MaxSim: each real query row's largest dot product with any real document row, summed; a `None` mask is all real.
/// An empty document scores `-1e9` and an empty query `0.0`; a mask of the wrong length or unequal widths panics.
pub fn late_interaction_score(
  query: &MultiVectorEmbedding,
  query_mask: Option<&[bool]>,
  document: &MultiVectorEmbedding,
  document_mask: Option<&[bool]>,
  masked_document_row_fill: MaskedRowFill,
) -> f32 {
  if let Some(mask) = query_mask {
    assert_eq!(
      mask.len(),
      query.row_count(),
      "query mask length must match its row count"
    );
  }
  if let Some(mask) = document_mask {
    assert_eq!(
      mask.len(),
      document.row_count(),
      "document mask length must match its row count"
    );
  }

  let document_is_real = |index: usize| document_mask.is_none_or(|mask| mask[index]);
  if !(0..document.row_count()).any(document_is_real) {
    return EMPTY_DOCUMENT_SCORE;
  }

  let query_is_real = |index: usize| query_mask.is_none_or(|mask| mask[index]);
  query
    .rows()
    .enumerate()
    .filter(|(index, _)| query_is_real(*index))
    .map(|(_, query_row)| {
      document
        .rows()
        .enumerate()
        .map(|(index, document_row)| {
          if document_is_real(index) {
            dot(query_row, document_row)
          } else {
            masked_document_row_fill.value()
          }
        })
        .fold(f32::MIN, f32::max)
    })
    .sum()
}

fn dot(a: &[f32], b: &[f32]) -> f32 {
  assert_eq!(
    a.len(),
    b.len(),
    "late interaction compares rows of equal width; got {} and {}",
    a.len(),
    b.len()
  );
  a.iter().zip(b).map(|(x, y)| x * y).sum()
}

mod reduction;

pub use reduction::ScoreReduction;
pub use reduction::late_interaction_score_for;

#[cfg(test)]
mod tests;