candding 0.1.0

//! `jinaai/jina-colbert-v2`'s head: `StanfordColbert`'s bias-free projection over the flash XLM-RoBERTa encoder,
//! whose queries attend to their mask-token padding when `attends_to_mask_tokens` says so.

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::multi_vector::l2_normalize_rows;
use crate::models::xlm_roberta_flash::XlmRobertaFlashConfig;
use crate::models::xlm_roberta_flash::XlmRobertaFlashModel;
use crate::multi_vector::MultiVectorEmbedding;

pub struct JinaColbert {
  encoder: XlmRobertaFlashModel,
  /// The root `linear.weight`, `(128, 1024)`, already `(out, in)`.
  weight: Tensor,
  /// The root `linear.bias` when `projection_bias` says so; `None` for this checkpoint.
  bias: Option<Tensor>,
  /// Ids of the 32 ASCII punctuation characters, dropped on the document side.
  punctuation_ids: HashSet<u32>,
  /// `weight`'s first dimension, the width of every returned row.
  dim: usize,
  normalize: bool,
  keeps_cls_row: bool,
  /// Whether a query's `[MASK]` fill is visible to the encoder's attention (`true` here); see `query_attention_mask`.
  attends_to_mask_tokens: bool,
}

impl JinaColbert {
  /// The projection sits at the root of the encoder's safetensors; `vb.get` checks its shape against the entry.
  pub fn load(
    config: XlmRobertaFlashConfig,
    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 = XlmRobertaFlashModel::load(config, vb, conventions.position_scheme)?;
    Ok(Self {
      encoder,
      weight,
      bias,
      punctuation_ids,
      dim: conventions.projection_out,
      normalize: conventions.normalize,
      keeps_cls_row: conventions.keeps_cls_row,
      attends_to_mask_tokens: conventions.attends_to_mask_tokens,
    })
  }

  /// `normalize(linear(hidden))` over every position, as in `StanfordColbert::project`.
  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.
  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(),
    )
  }

  /// The query's encoder attention mask: all ones when `attends_to_mask_tokens`, else unchanged.
  fn query_attention_mask(&self, attention_mask: &Tensor) -> Result<Tensor> {
    if self.attends_to_mask_tokens {
      Ok(Tensor::ones(
        attention_mask.shape(),
        attention_mask.dtype(),
        attention_mask.device(),
      )?)
    } else {
      Ok(attention_mask.clone())
    }
  }
}

impl MultiVectorModel for JinaColbert {
  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 mask = self.query_attention_mask(input.attention_mask)?;
    let effective = EncoderInput {
      input_ids: input.input_ids,
      attention_mask: &mask,
      token_type_ids: input.token_type_ids,
      lengths: input.lengths,
    };
    let hidden = self.encoder.forward(&effective)?;
    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;