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::DtypeSupport;
use super::Family;
use super::Support;
use super::WeightsFormat;
use super::entry::Entry;
use super::t4::T4_BF16_UNTESTED;
use crate::pooling::Pooling;

/// Neither e5 repository ships prompts; the model card gives `query: ` for symmetric tasks too.
const QUERY: &str = "query: {text}";
const PASSAGE: &str = "passage: {text}";

/// A multilingual-e5 dense entry: mean pooling and the `query: ` and `passage: ` prefixes.
const fn e5(id: &'static str, dim: usize, size_gb: f32) -> Entry {
  Entry {
    id,
    family: Family::XlmRoberta,
    dim,
    default_max_length: 512,
    pooling: Pooling::Mean,
    cpu: Support::Verified,
    metal: Support::Verified,
    cuda: Support::Verified,
    query_template: Some(QUERY),
    passage_template: Some(PASSAGE),
    default_instruction: None,
    dense_dirs: &[],
    mrl: false,
    layer_norm_after_pooling: false,
    gated: false,
    default_dtype: DType::F32,
    dtypes: DtypeSupport::verified(),
    license: "MIT",
    size_gb,
    phase: 2,
    weights: WeightsFormat::Safetensors,
    accel_only_dtypes: &[],
    cell_exceptions: T4_BF16_UNTESTED,
  }
}

/// `size_gb` is the decimal-gigabyte download size of `model.safetensors`: 1 112 201 288 and 2 239 611 368 bytes.
pub(super) const ENTRIES: &[Entry] = &[
  e5("intfloat/multilingual-e5-base", 768, 1.11),
  e5("intfloat/multilingual-e5-large", 1024, 2.24),
  BGE_M3,
  PARAPHRASE_MULTILINGUAL_MPNET,
];

/// `BAAI/bge-m3`: a `pytorch_model.bin` and no safetensors, so the fetch adds the pickles, the sparse head's
/// `sparse_linear.pt` and the multi-vector head's `colbert_linear.pt` included, which load through their own entries.
const BGE_M3: Entry = Entry {
  id: "BAAI/bge-m3",
  family: Family::XlmRoberta,
  dim: 1024,
  default_max_length: 8192,
  pooling: Pooling::Cls,
  cpu: Support::Verified,
  metal: Support::Verified,
  cuda: Support::Verified,
  query_template: None,
  passage_template: None,
  default_instruction: None,
  dense_dirs: &[],
  mrl: false,
  layer_norm_after_pooling: false,
  gated: false,
  default_dtype: DType::F32,
  dtypes: DtypeSupport::verified(),
  license: "MIT",
  size_gb: 2.27,
  phase: 2,
  weights: WeightsFormat::Pickle,
  accel_only_dtypes: &[],
  cell_exceptions: T4_BF16_UNTESTED,
};

/// `sentence-transformers/paraphrase-multilingual-mpnet-base-v2` is an XLM-RoBERTa by its `config.json`, despite
/// its name; it ships no `Normalize` module, and its 128-token default is its `sentence_bert_config.json` value.
const PARAPHRASE_MULTILINGUAL_MPNET: Entry = Entry {
  id: "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
  family: Family::XlmRoberta,
  dim: 768,
  default_max_length: 128,
  pooling: Pooling::Mean,
  cpu: Support::Verified,
  metal: Support::Verified,
  cuda: Support::Verified,
  query_template: None,
  passage_template: None,
  default_instruction: None,
  dense_dirs: &[],
  mrl: false,
  layer_norm_after_pooling: false,
  gated: false,
  default_dtype: DType::F32,
  dtypes: DtypeSupport::verified(),
  license: "Apache-2.0",
  size_gb: 1.11,
  phase: 2,
  weights: WeightsFormat::Safetensors,
  accel_only_dtypes: &[],
  cell_exceptions: T4_BF16_UNTESTED,
};