candding 0.1.0

//! The multi-vector trait and its dispatch, a third output kind beside `Encoder`'s dense tensor and `SparseModel`'s
//! rows; `load_multi_vector` builds the head a `MultiVectorHead` names.

mod jina_colbert;
mod m3_colbert;
mod stanford_colbert;

use std::path::Path;
use std::path::PathBuf;

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

use super::EncoderInput;
use crate::error::CanddingError;
use crate::error::Result;
use crate::loader::COLBERT_LINEAR_FILENAME;
use crate::multi_vector::MultiVectorEmbedding;

/// L2-normalizes every row of a tensor whose last dimension is the row width, clamping the norm at `1e-12` as
/// `embed::postprocess::normalize` does; `models` cannot depend on `embed`, so that function is not reused.
pub(super) fn l2_normalize_rows(xs: &Tensor) -> Result<Tensor> {
  let last = xs.rank() - 1;
  let norms = xs.sqr()?.sum_keepdim(last)?.sqrt()?.maximum(1e-12)?;
  Ok(xs.broadcast_div(&norms)?)
}

/// A multi-vector head: one `MultiVectorEmbedding` per row, projection and normalization included. `forward` and
/// `forward_query` share no default, since every checkpoint marks or pads a query differently from a document.
pub trait MultiVectorModel: Send + Sync {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>>;
  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>>;
  /// The projected row width every returned `MultiVectorEmbedding` carries, not the encoder's hidden size.
  fn dim(&self) -> usize;
  fn max_position_embeddings(&self) -> usize;
  /// `Encoder::weights_dtype` of the encoder the head runs on.
  fn weights_dtype(&self) -> DType;
}

/// Which multi-vector head a checkpoint runs.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum MultiVectorHead {
  /// A biased `Linear(1024, 1024)` that drops the classification row: `BAAI/bge-m3`.
  M3Colbert,
  /// A bias-free projection that keeps the classification row: `colbert-ir/colbertv2.0` and
  /// `answerdotai/answerai-colbert-small-v1`.
  StanfordColbert,
  /// The same projection over the flash XLM-RoBERTa encoder, the query attending to its mask-token padding:
  /// `jinaai/jina-colbert-v2`.
  JinaColbert,
}

/// The per-checkpoint facts `load_multi_vector` reads; a `models`-side copy of the registry entry's fields,
/// since `models` cannot depend on `registry`.
#[derive(Clone, Copy, Debug)]
pub struct MultiVectorConventions {
  pub keeps_cls_row: bool,
  pub projection_in: usize,
  pub projection_out: usize,
  pub projection_bias: bool,
  pub normalize: bool,
  /// Whether a query's mask-token padding is attended to (`true` for jina-colbert-v2); read by `JinaColbert`.
  pub attends_to_mask_tokens: bool,
  /// The flash XLM-RoBERTa position scheme (`Some(Rotary)` for jina-colbert-v2, else `None`), checked at load.
  pub position_scheme: Option<super::xlm_roberta_flash::PositionScheme>,
}

/// Builds the multi-vector head `head` names, parallel to `load_sparse`. Only `M3Colbert` reads `colbert_linear`
/// and only `StanfordColbert` reads `tokenizer_path`, for its punctuation skiplist.
pub fn load_multi_vector(
  head: MultiVectorHead,
  config_json: &str,
  vb: VarBuilder,
  colbert_linear: Option<&Path>,
  tokenizer_path: &Path,
  conventions: MultiVectorConventions,
) -> Result<Box<dyn MultiVectorModel>> {
  match head {
    MultiVectorHead::M3Colbert => {
      let colbert_linear = colbert_linear
        .ok_or_else(|| CanddingError::MissingFile(PathBuf::from(COLBERT_LINEAR_FILENAME)))?;
      Ok(Box::new(m3_colbert::M3Colbert::load(
        serde_json::from_str(config_json)?,
        vb,
        colbert_linear,
        conventions,
      )?))
    }
    MultiVectorHead::StanfordColbert => {
      let punctuation_ids = crate::tokenizer::ascii_punctuation_ids(tokenizer_path)?;
      Ok(Box::new(stanford_colbert::StanfordColbert::load(
        serde_json::from_str(config_json)?,
        vb,
        punctuation_ids,
        conventions,
      )?))
    }
    MultiVectorHead::JinaColbert => {
      let punctuation_ids = crate::tokenizer::ascii_punctuation_ids(tokenizer_path)?;
      Ok(Box::new(jina_colbert::JinaColbert::load(
        serde_json::from_str(config_json)?,
        vb,
        punctuation_ids,
        conventions,
      )?))
    }
  }
}

#[cfg(test)]
mod tests;