candding 0.1.0

//! Multi-vector text embeddings: one `MultiVectorEmbedding` (many rows, not one pooled vector) per text, over a
//! model's own multi-vector head, parallel to `SparseTextEmbedding` in the sibling module.

mod build;
mod builder;

pub use builder::MultiVectorTextEmbeddingBuilder;
use candle_core::DType;
use candle_core::Device;

use super::METAL_FORWARD_LOCK;
use super::batching;
use crate::error::Result;
use crate::models::EncoderInput;
use crate::models::MultiVectorModel;
use crate::multi_vector::MultiVectorEmbedding;
use crate::registry::multi_vector::MultiVectorDescriptor;
use crate::tokenizer::TextTokenizer;

/// Which side of a late-interaction pair a batch is embedded for. It picks the marker, the tokenizer and the
/// forward method, because a fixed-length query is padded at tokenization time.
#[derive(Clone, Copy)]
enum Role {
  Query,
  Document,
}

/// A loaded multi-vector model: one `MultiVectorEmbedding` of per-token rows per text. A checkpoint that pads
/// queries to a fixed length gets a query-side tokenizer of its own, and Metal forward passes share one lock.
pub struct MultiVectorTextEmbedding {
  model: Box<dyn MultiVectorModel>,
  tokenizer: TextTokenizer,
  query_tokenizer: TextTokenizer,
  descriptor: MultiVectorDescriptor,
  device: Device,
  dtype: DType,
  batch_size: usize,
}

impl MultiVectorTextEmbedding {
  pub fn builder(spec: impl Into<String>) -> MultiVectorTextEmbeddingBuilder {
    MultiVectorTextEmbeddingBuilder::new(spec)
  }

  #[allow(clippy::too_many_arguments)]
  pub(crate) fn new(
    model: Box<dyn MultiVectorModel>,
    tokenizer: TextTokenizer,
    query_tokenizer: TextTokenizer,
    descriptor: MultiVectorDescriptor,
    device: Device,
    dtype: DType,
    batch_size: usize,
  ) -> Self {
    Self {
      model,
      tokenizer,
      query_tokenizer,
      descriptor,
      device,
      dtype,
      batch_size,
    }
  }

  pub fn descriptor(&self) -> &MultiVectorDescriptor {
    &self.descriptor
  }

  pub fn max_length(&self) -> usize {
    self.tokenizer.max_length()
  }

  pub fn device(&self) -> &Device {
    &self.device
  }

  pub fn dtype(&self) -> DType {
    self.dtype
  }

  /// The dtype the loaded weights are in, read back from the encoder's token embedding table
  /// (`MultiVectorModel::weights_dtype`); `dtype` is the one the builder resolved and asked the loader for.
  pub fn weights_dtype(&self) -> DType {
    self.model.weights_dtype()
  }

  /// Token counts per text after truncation, special tokens included, through the document-side tokenizer and
  /// without a query or document marker.
  pub fn token_count<S: AsRef<str>>(&self, texts: &[S]) -> Result<Vec<usize>> {
    let refs: Vec<&str> = texts.iter().map(|t| t.as_ref()).collect();
    self.tokenizer.count(&refs)
  }

  /// The document side, the same rows as `passage_embed`. Batches of `batch_size` (the builder's value when
  /// `None`) are sorted longest-first, then restored to the caller's order.
  pub fn embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<MultiVectorEmbedding>> {
    self.passage_embed(texts, batch_size)
  }

  /// The document side: the document marker, where the checkpoint has one, and the document tokenizer.
  pub fn passage_embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<MultiVectorEmbedding>> {
    self.embed_with(texts, batch_size, Role::Document)
  }

  /// The query side: the query marker and, for a checkpoint with a fixed query length, its `[MASK]`-padded
  /// tokenizer; the passage rows for `BAAI/bge-m3`, which has neither.
  pub fn query_embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<MultiVectorEmbedding>> {
    self.embed_with(texts, batch_size, Role::Query)
  }

  fn embed_with<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
    role: Role,
  ) -> Result<Vec<MultiVectorEmbedding>> {
    if texts.is_empty() {
      return Ok(Vec::new());
    }
    let batch_size = batch_size.unwrap_or(self.batch_size).max(1);
    let formatted = self.formatted(role, texts);
    let refs: Vec<&str> = formatted.iter().map(String::as_str).collect();
    let tokenizer = self.tokenizer_for(role);
    let lengths = tokenizer.count(&refs)?;
    let order = batching::sorted_indices(&lengths);
    let mut rows = Vec::with_capacity(texts.len());
    for chunk in order.chunks(batch_size) {
      let chunk_texts: Vec<&str> = chunk.iter().map(|i| refs[*i]).collect();
      rows.extend(self.embed_chunk(&chunk_texts, tokenizer, role)?);
    }
    Ok(batching::restore(&order, rows))
  }

  /// Prepends `role`'s marker by plain concatenation: the marker already ends in the reference prompt's space
  /// (`"[unused0] "`), so none is added here.
  fn formatted<S: AsRef<str>>(&self, role: Role, texts: &[S]) -> Vec<String> {
    let marker = match role {
      Role::Query => self.descriptor.query_marker.as_deref(),
      Role::Document => self.descriptor.document_marker.as_deref(),
    };
    texts
      .iter()
      .map(|t| match marker {
        Some(marker) => format!("{marker}{}", t.as_ref()),
        None => t.as_ref().to_string(),
      })
      .collect()
  }

  fn tokenizer_for(&self, role: Role) -> &TextTokenizer {
    match role {
      Role::Query => &self.query_tokenizer,
      Role::Document => &self.tokenizer,
    }
  }

  /// As in `SparseTextEmbedding::embed_chunk`, an all-empty chunk skips `encode`, since Metal refuses empty buffers.
  fn embed_chunk(
    &self,
    texts: &[&str],
    tokenizer: &TextTokenizer,
    role: Role,
  ) -> Result<Vec<MultiVectorEmbedding>> {
    if tokenizer.count(texts)?.iter().all(|&count| count == 0) {
      let empty = MultiVectorEmbedding::from_rows(self.model.dim(), Vec::<Vec<f32>>::new());
      return Ok(vec![empty; texts.len()]);
    }
    let _guard = self.device.is_metal().then(|| {
      METAL_FORWARD_LOCK
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
    });
    let batch = tokenizer.encode(texts, &self.device)?;
    let input = EncoderInput {
      input_ids: &batch.input_ids,
      attention_mask: &batch.attention_mask,
      token_type_ids: &batch.token_type_ids,
      lengths: &batch.lengths,
    };
    match role {
      Role::Query => self.model.forward_query(&input),
      Role::Document => self.model.forward(&input),
    }
  }
}

#[cfg(test)]
mod tests;