candding 0.1.0

//! Sparse text embeddings: a row of `(vocabulary index, weight)` pairs per text, over a model's own sparse
//! head, in place of a fixed-width dense vector.

mod build;
mod builder;

pub use builder::SparseTextEmbeddingBuilder;
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::SparseModel;
use crate::registry::sparse::SparseDescriptor;
use crate::sparse::SparseEmbedding;
use crate::tokenizer::TextTokenizer;

/// A loaded sparse model that runs an encoder: one `SparseEmbedding` per text. On Metal its forward passes share
/// the process-wide lock `TextEmbedding` takes.
pub struct SparseTextEmbedding {
  model: Box<dyn SparseModel>,
  tokenizer: TextTokenizer,
  descriptor: SparseDescriptor,
  device: Device,
  dtype: DType,
  batch_size: usize,
}

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

  pub(crate) fn new(
    model: Box<dyn SparseModel>,
    tokenizer: TextTokenizer,
    descriptor: SparseDescriptor,
    device: Device,
    dtype: DType,
    batch_size: usize,
  ) -> Self {
    Self {
      model,
      tokenizer,
      descriptor,
      device,
      dtype,
      batch_size,
    }
  }

  pub fn descriptor(&self) -> &SparseDescriptor {
    &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
  /// (`SparseModel::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.
  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 passage 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<SparseEmbedding>> {
    self.passage_embed(texts, batch_size)
  }

  /// The passage side, through `SparseModel::forward`.
  pub fn passage_embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<SparseEmbedding>> {
    self.embed_with(texts, batch_size, |model, input| model.forward(input))
  }

  /// The query side, through `SparseModel::forward_query`: the passage rows for every head but `MiniCoil` and
  /// `Bm42`, which weigh a query's terms uniformly.
  pub fn query_embed<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<SparseEmbedding>> {
    self.embed_with(texts, batch_size, |model, input| model.forward_query(input))
  }

  fn embed_with<S: AsRef<str>>(
    &self,
    texts: &[S],
    batch_size: Option<usize>,
    forward: fn(&dyn SparseModel, &EncoderInput) -> Result<Vec<SparseEmbedding>>,
  ) -> Result<Vec<SparseEmbedding>> {
    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(|t| t.as_ref()).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, forward)?);
    }
    Ok(batching::restore(&order, rows))
  }

  /// A chunk whose every text tokenizes to zero tokens (splade-code's empty string) never reaches `encode`:
  /// its `(batch, 0)` tensor would be a zero-byte buffer, which candle's Metal backend refuses to allocate.
  fn embed_chunk(
    &self,
    texts: &[&str],
    forward: fn(&dyn SparseModel, &EncoderInput) -> Result<Vec<SparseEmbedding>>,
  ) -> Result<Vec<SparseEmbedding>> {
    if self.tokenizer.count(texts)?.iter().all(|&count| count == 0) {
      return Ok(vec![SparseEmbedding::from_pairs([]); texts.len()]);
    }
    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,
    };
    forward(self.model.as_ref(), &input)
  }
}

#[cfg(test)]
mod tests;