candding 0.1.0

//! Dense text embeddings.

use candle_core::DType;
use candle_core::Device;

use super::METAL_FORWARD_LOCK;
use super::batching;
use super::builder::TextEmbeddingBuilder;
use super::postprocess;
use crate::error::CanddingError;
use crate::error::Result;
use crate::layers::dense::Dense;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::pooling;
use crate::registry::ModelDescriptor;
use crate::template;
use crate::tokenizer::TextTokenizer;

/// A loaded dense model. On Metal every forward pass in the process is serialized behind one lock; on the CPU,
/// calls from separate threads run concurrently.
pub struct TextEmbedding {
  encoder: Box<dyn Encoder>,
  tokenizer: TextTokenizer,
  dense: Vec<Dense>,
  descriptor: ModelDescriptor,
  device: Device,
  dtype: DType,
  batch_size: usize,
}

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

  pub(crate) fn new(
    encoder: Box<dyn Encoder>,
    tokenizer: TextTokenizer,
    dense: Vec<Dense>,
    descriptor: ModelDescriptor,
    device: Device,
    dtype: DType,
    batch_size: usize,
  ) -> Self {
    Self {
      encoder,
      tokenizer,
      dense,
      descriptor,
      device,
      dtype,
      batch_size,
    }
  }

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

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

  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
  /// (`Encoder::weights_dtype`); `dtype` is the one the builder resolved and asked the loader for.
  pub fn weights_dtype(&self) -> DType {
    self.encoder.weights_dtype()
  }

  /// Embeds texts with no role applied: neither the query nor the passage template, unlike `embed` on the
  /// sparse, multi-vector and BM25 types, where it is the passage side.
  pub fn embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(plain(texts), None, batch_size)
  }

  pub fn embed_dim<S: AsRef<str>>(
    &self,
    texts: &[S],
    dim: usize,
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(plain(texts), Some(dim), batch_size)
  }

  pub fn query_embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(
      self.queries(self.descriptor.default_instruction.as_deref(), texts),
      None,
      batch_size,
    )
  }

  pub fn query_embed_dim<S: AsRef<str>>(
    &self,
    texts: &[S],
    dim: usize,
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(
      self.queries(self.descriptor.default_instruction.as_deref(), texts),
      Some(dim),
      batch_size,
    )
  }

  /// Replaces the descriptor's default instruction for instruct-style query templates; the instruction comes
  /// after the texts, as in `TextCrossEncoder::rerank_with_instruction`.
  pub fn query_embed_with_instruction<S: AsRef<str>>(
    &self,
    texts: &[S],
    instruction: &str,
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(self.queries(Some(instruction), texts), None, batch_size)
  }

  pub fn passage_embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(self.passages(texts), None, batch_size)
  }

  pub fn passage_embed_dim<S: AsRef<str>>(
    &self,
    texts: &[S],
    dim: usize,
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    self.run(self.passages(texts), Some(dim), batch_size)
  }

  /// Token counts per text after truncation, special tokens included, without templates.
  pub fn token_count<S: AsRef<str>>(&self, texts: &[S]) -> Result<Vec<usize>> {
    self.count(plain(texts))
  }

  /// The counts `query_embed` would encode: the descriptor's query template around each text, with its default
  /// instruction. A template is part of the input the model sees, so it is part of what a count of that run means.
  pub fn query_token_count<S: AsRef<str>>(&self, texts: &[S]) -> Result<Vec<usize>> {
    self.count(self.queries(self.descriptor.default_instruction.as_deref(), texts))
  }

  /// The counts `query_embed_with_instruction` would encode: the query template around each text with the caller's
  /// instruction in place of the descriptor's default.
  pub fn query_token_count_with_instruction<S: AsRef<str>>(
    &self,
    texts: &[S],
    instruction: &str,
  ) -> Result<Vec<usize>> {
    self.count(self.queries(Some(instruction), texts))
  }

  /// The counts `passage_embed` would encode.
  pub fn passage_token_count<S: AsRef<str>>(&self, texts: &[S]) -> Result<Vec<usize>> {
    self.count(self.passages(texts))
  }

  fn count(&self, texts: Vec<String>) -> Result<Vec<usize>> {
    let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
    self.tokenizer.count(&refs)
  }

  fn queries<S: AsRef<str>>(&self, instruction: Option<&str>, texts: &[S]) -> Vec<String> {
    let template = self.descriptor.query_template.as_deref();
    texts
      .iter()
      .map(|t| template::apply(template, instruction, t.as_ref()))
      .collect()
  }

  fn passages<S: AsRef<str>>(&self, texts: &[S]) -> Vec<String> {
    let template = self.descriptor.passage_template.as_deref();
    texts
      .iter()
      .map(|t| template::apply(template, None, t.as_ref()))
      .collect()
  }

  fn run(
    &self,
    texts: Vec<String>,
    dim: Option<usize>,
    batch_size: Option<usize>,
  ) -> Result<Vec<Vec<f32>>> {
    if dim.is_some() && !self.descriptor.mrl {
      return Err(CanddingError::MrlUnsupported(self.descriptor.id.clone()));
    }
    if texts.is_empty() {
      return Ok(Vec::new());
    }
    let batch_size = batch_size.unwrap_or(self.batch_size).max(1);
    let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
    let lengths = self.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, dim)?);
    }
    Ok(batching::restore(&order, rows))
  }

  /// Pools, upcasts to F32, applies Dense, the optional layer norm, truncation and normalization; on Metal the
  /// whole body holds `METAL_FORWARD_LOCK`, since every op allocates from the shared buffer pool.
  fn embed_chunk(&self, texts: &[&str], dim: Option<usize>) -> Result<Vec<Vec<f32>>> {
    let _guard = self.device.is_metal().then(|| {
      METAL_FORWARD_LOCK
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
    });
    let batch = self.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,
    };
    let hidden = self.encoder.forward(&input)?;
    let pooled = pooling::pool(
      self.descriptor.pooling,
      &hidden,
      &batch.attention_mask,
      &batch.lengths,
    )?
    .to_dtype(DType::F32)?;
    let pooled = self
      .dense
      .iter()
      .try_fold(pooled, |xs, dense| dense.forward(&xs))?;
    let pooled = if self.descriptor.layer_norm_after_pooling {
      postprocess::layer_norm(&pooled, 1e-5)?
    } else {
      pooled
    };
    let pooled = match dim {
      Some(dim) => postprocess::truncate_dim(&pooled, dim)?,
      None => pooled,
    };
    let pooled = if self.descriptor.normalize {
      postprocess::normalize(&pooled)?
    } else {
      pooled
    };
    postprocess::to_rows(&pooled)
  }
}

fn plain<S: AsRef<str>>(texts: &[S]) -> Vec<String> {
  texts.iter().map(|t| t.as_ref().to_string()).collect()
}

#[cfg(test)]
mod tests;