candding 0.1.0

//! Shared head for `colbert-ir/colbertv2.0` and `answerdotai/answerai-colbert-small-v1`: a bias-free root projection
//! of every position of a BERT encoder, classification row included, with their shared markers, padding and skiplist.

use std::collections::HashSet;

use candle_core::DType;
use candle_core::Tensor;
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::bert::BertConfig;
use crate::models::bert::BertModel;
use crate::models::multi_vector::l2_normalize_rows;
use crate::multi_vector::MultiVectorEmbedding;

pub struct StanfordColbert {
  encoder: BertModel,
  /// The root `linear.weight`, `(projection_out, hidden_size)`, already `(out, in)`.
  weight: Tensor,
  /// The root `linear.bias` when `projection_bias` says so; `None` for both registered checkpoints.
  bias: Option<Tensor>,
  /// Ids of the 32 ASCII punctuation characters, the document-side skiplist; queries keep them, as in the reference.
  punctuation_ids: HashSet<u32>,
  /// `weight`'s first dimension, the width of every returned row.
  dim: usize,
  normalize: bool,
  /// Whether position 0 survives; `true` for both checkpoints, unlike bge-m3.
  keeps_cls_row: bool,
}

impl StanfordColbert {
  /// The projection sits at the root of the encoder's safetensors; `vb.get` checks its shape against the entry.
  pub fn load(
    config: BertConfig,
    vb: VarBuilder,
    punctuation_ids: HashSet<u32>,
    conventions: MultiVectorConventions,
  ) -> Result<Self> {
    if conventions.projection_in != config.hidden_size {
      return Err(CanddingError::ProjectionInMismatch {
        declared: conventions.projection_in,
        actual: config.hidden_size,
      });
    }
    let weight = vb.get(
      (conventions.projection_out, config.hidden_size),
      "linear.weight",
    )?;
    let bias = if conventions.projection_bias {
      Some(vb.get(conventions.projection_out, "linear.bias")?)
    } else {
      None
    };
    let encoder = BertModel::load(config, vb)?;
    Ok(Self {
      encoder,
      weight,
      bias,
      punctuation_ids,
      dim: conventions.projection_out,
      normalize: conventions.normalize,
      keeps_cls_row: conventions.keeps_cls_row,
    })
  }

  /// `normalize(linear(hidden))` over every position, before either row selection.
  fn project(&self, hidden: &Tensor) -> Result<Vec<Vec<Vec<f32>>>> {
    let linear = Linear::new(self.weight.clone(), self.bias.clone());
    let projected_raw = linear.forward(hidden)?;
    let projected = if self.normalize {
      l2_normalize_rows(&projected_raw)?
    } else {
      projected_raw
    };
    Ok(projected.to_dtype(DType::F32)?.to_vec3::<f32>()?)
  }

  /// Every row except padding, document-side punctuation and, unless `keeps_cls_row`, the first; a masked gather
  /// since punctuation sits anywhere.
  fn select_document_rows(
    &self,
    hidden: &Tensor,
    ids: &[Vec<u32>],
    lengths: &[usize],
  ) -> Result<Vec<MultiVectorEmbedding>> {
    let batch = self.project(hidden)?;
    Ok(
      batch
        .into_iter()
        .zip(ids.iter().cloned())
        .zip(lengths)
        .map(|((positions, row_ids), &real_len)| {
          let kept = positions
            .into_iter()
            .zip(row_ids)
            .enumerate()
            .filter(|(position, (_, id))| {
              *position < real_len
                && !self.punctuation_ids.contains(id)
                && (self.keeps_cls_row || *position != 0)
            })
            .map(|(_, (row, _))| row);
          MultiVectorEmbedding::from_rows(self.dim, kept)
        })
        .collect(),
    )
  }

  /// Every position: the reference scores the whole fixed-width query, `[MASK]` fill included.
  fn select_query_rows(&self, hidden: &Tensor) -> Result<Vec<MultiVectorEmbedding>> {
    let batch = self.project(hidden)?;
    Ok(
      batch
        .into_iter()
        .map(|positions| MultiVectorEmbedding::from_rows(self.dim, positions))
        .collect(),
    )
  }
}

impl MultiVectorModel for StanfordColbert {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    let hidden = self.encoder.forward(input)?;
    let ids = input.input_ids.to_vec2::<u32>()?;
    self.select_document_rows(&hidden, &ids, input.lengths)
  }

  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    let hidden = self.encoder.forward(input)?;
    self.select_query_rows(&hidden)
  }

  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;