candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
Documentation
use candle_core::DType;

use super::BackendSupport;
use super::CellException;
use super::DtypeBackendSupport;
use super::DtypeSupport;
use super::Family;
use super::ModelDescriptor;
use super::Support;
use super::WeightsFormat;
use crate::pooling::Pooling;

/// Compact static form of a descriptor.
pub(crate) struct Entry {
  pub id: &'static str,
  pub family: Family,
  pub dim: usize,
  pub default_max_length: usize,
  pub pooling: Pooling,
  pub cpu: Support,
  pub metal: Support,
  pub cuda: Support,
  pub query_template: Option<&'static str>,
  pub passage_template: Option<&'static str>,
  pub default_instruction: Option<&'static str>,
  pub dense_dirs: &'static [&'static str],
  pub mrl: bool,
  /// A parameter-free LayerNorm over the pooled vector, after Dense and before Matryoshka truncation.
  pub layer_norm_after_pooling: bool,
  pub gated: bool,
  pub default_dtype: DType,
  pub dtypes: DtypeSupport,
  pub license: &'static str,
  pub size_gb: f32,
  pub phase: u8,
  /// Off `ModelDescriptor` on purpose: the schema and site types do not change for one model.
  pub weights: WeightsFormat,
  /// Dtypes correct on an accelerator but not on the CPU (candle's CPU f16 matmul accumulates in f16), which
  /// `DtypeBackendSupport::derive` turns into a `cpu: Unsupported` row the builder refuses.
  pub accel_only_dtypes: &'static [DType],
  /// Cells a run decided that the derivation cannot state, written over its rows (`CellException`).
  pub cell_exceptions: &'static [CellException],
}

impl Entry {
  pub fn descriptor(&self) -> ModelDescriptor {
    let backends = BackendSupport {
      cpu: self.cpu,
      metal: self.metal,
      cuda: self.cuda,
    };
    ModelDescriptor {
      id: self.id.to_string(),
      family: self.family,
      dim: self.dim,
      default_max_length: self.default_max_length,
      pooling: self.pooling,
      normalize: true,
      query_template: self.query_template.map(str::to_string),
      passage_template: self.passage_template.map(str::to_string),
      default_instruction: self.default_instruction.map(str::to_string),
      dense_dirs: self.dense_dirs.iter().map(|d| d.to_string()).collect(),
      mrl: self.mrl,
      layer_norm_after_pooling: self.layer_norm_after_pooling,
      license: self.license.to_string(),
      size_gb: self.size_gb,
      gated: self.gated,
      phase: self.phase,
      backends,
      default_dtype: self.default_dtype,
      dtype_backends: DtypeBackendSupport::derive(backends, self.dtypes, self.accel_only_dtypes)
        .with_exceptions(self.cell_exceptions),
    }
  }

  /// This entry as `support::rules` reads it.
  #[cfg(test)]
  pub(crate) fn cells(&self) -> super::support::rules::Cells<'static> {
    let descriptor = self.descriptor();
    super::support::rules::Cells {
      id: self.id,
      backends: descriptor.backends,
      dtypes: self.dtypes,
      accel_only: self.accel_only_dtypes,
      exceptions: self.cell_exceptions,
      rows: descriptor.dtype_backends,
    }
  }
}