candding 0.1.0

//! `BAAI/bge-m3`'s multi-vector head: `colbert_linear.pt`'s biased `Linear(1024, 1024)` on every position after the
//! classification row, which it drops; rows are L2-normalized and trimmed to each input's real length.

use std::path::Path;

use candle_core::DType;
use candle_core::Tensor;
use candle_core::pickle::PthTensors;
use candle_nn::Linear;
use candle_nn::Module;
use candle_nn::VarBuilder;

use crate::error::CanddingError;
use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::MultiVectorConventions;
use crate::models::MultiVectorModel;
use crate::models::multi_vector::l2_normalize_rows;
use crate::models::xlm_roberta::XlmRobertaConfig;
use crate::models::xlm_roberta::XlmRobertaModel;
use crate::multi_vector::MultiVectorEmbedding;

/// `BAAI/bge-m3`'s multi-vector head, built on `xlm_roberta::XlmRobertaModel`.
#[derive(Debug)]
pub struct M3Colbert {
  encoder: XlmRobertaModel,
  /// `colbert_linear.pt`'s `(1024, 1024)` weight, already `(out, in)`.
  weight: Tensor,
  /// `colbert_linear.pt`'s `(1024,)` bias, when `projection_bias` says the checkpoint has one.
  bias: Option<Tensor>,
  /// `weight`'s first dimension, the width of every returned row.
  dim: usize,
  normalize: bool,
  /// Whether position 0 survives; `false` for bge-m3, unlike the Stanford-ColBERT checkpoints.
  keeps_cls_row: bool,
}

/// Like `PthTensors::get`, but a missing tensor names the file too; `sparse::m3`'s copy is private.
fn read_tensor(pth: &PthTensors, name: &str, file: &Path) -> Result<Tensor> {
  pth.get(name)?.ok_or_else(|| CanddingError::MissingTensor {
    file: file.to_path_buf(),
    name: name.to_string(),
  })
}

impl M3Colbert {
  /// Reads `colbert_linear.pt`'s float16 weight, and its bias when `conventions.projection_bias` says so; a missing
  /// file errors by name, where FlagEmbedding's loader would build a randomly initialized head.
  pub fn load(
    config: XlmRobertaConfig,
    vb: VarBuilder,
    colbert_linear: &Path,
    conventions: MultiVectorConventions,
  ) -> Result<Self> {
    if !colbert_linear.is_file() {
      return Err(CanddingError::MissingFile(colbert_linear.to_path_buf()));
    }
    let pth = PthTensors::new(colbert_linear, None)?;
    let weight = read_tensor(&pth, "weight", colbert_linear)?;
    let bias = if conventions.projection_bias {
      Some(read_tensor(&pth, "bias", colbert_linear)?)
    } else {
      None
    };
    Self::from_raw_head(config, vb, weight, bias, conventions)
  }

  /// Checks `projection_in` against the weight, casts the head and builds the encoder; split out for hand-built tests.
  fn from_raw_head(
    config: XlmRobertaConfig,
    vb: VarBuilder,
    weight: Tensor,
    bias: Option<Tensor>,
    conventions: MultiVectorConventions,
  ) -> Result<Self> {
    let actual_in = weight.dim(1)?;
    if actual_in != conventions.projection_in {
      return Err(CanddingError::ProjectionInMismatch {
        declared: conventions.projection_in,
        actual: actual_in,
      });
    }
    let dim = weight.dim(0)?;
    let dtype = vb.dtype();
    let weight = weight.to_dtype(dtype)?.to_device(vb.device())?;
    let bias = bias
      .map(|b| -> Result<Tensor> { Ok(b.to_dtype(dtype)?.to_device(vb.device())?) })
      .transpose()?;
    let encoder = XlmRobertaModel::load(config, vb)?;
    Ok(Self {
      encoder,
      weight,
      bias,
      dim,
      normalize: conventions.normalize,
      keeps_cls_row: conventions.keeps_cls_row,
    })
  }

  /// `normalize(linear(hidden))` minus position 0 unless `keeps_cls_row`, trimmed to each row's real length;
  /// `seq.min(drop)` keeps an empty batch from underflowing.
  fn project_rows(&self, hidden: &Tensor, lengths: &[usize]) -> Result<Vec<MultiVectorEmbedding>> {
    let seq = hidden.dim(1)?;
    let drop = if self.keeps_cls_row { 0 } else { 1 };
    let width = seq.saturating_sub(drop);
    let narrowed = hidden.narrow(1, seq.min(drop), width)?;
    let linear = Linear::new(self.weight.clone(), self.bias.clone());
    let projected_raw = linear.forward(&narrowed)?;
    let projected = if self.normalize {
      l2_normalize_rows(&projected_raw)?
    } else {
      projected_raw
    };
    let batch = projected
      .to_dtype(candle_core::DType::F32)?
      .to_vec3::<f32>()?;
    Ok(
      batch
        .into_iter()
        .zip(lengths)
        .map(|(rows, &len)| {
          let real_rows = len.saturating_sub(drop).min(rows.len());
          MultiVectorEmbedding::from_rows(self.dim, rows.into_iter().take(real_rows))
        })
        .collect(),
    )
  }
}

impl MultiVectorModel for M3Colbert {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    let hidden = self.encoder.forward(input)?;
    self.project_rows(&hidden, input.lengths)
  }

  /// Identical to `forward`: this checkpoint has no query padding or marker step.
  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    self.forward(input)
  }

  fn dim(&self) -> usize {
    self.dim
  }

  fn max_position_embeddings(&self) -> usize {
    self.encoder.max_position_embeddings()
  }

  fn weights_dtype(&self) -> DType {
    self.encoder.weights_dtype()
  }
}

#[cfg(test)]
mod tests;