candding 0.1.0

//! Built-in model descriptors.

mod bert;
mod catalog;
pub mod dense;
mod entry;
mod gemma3;
mod gte_new;
mod jina_bert;
mod modernbert;
mod mpnet;
pub mod multi_vector;
mod nomic_bert;
mod qwen3;
pub mod rerank;
pub mod sparse;
mod support;
mod t4;
mod xlm_roberta;

use candle_core::DType;
pub use catalog::Catalog;
pub(crate) use catalog::other_catalog;
use entry::Entry;
use serde::Deserialize;
use serde::Serialize;
pub use support::Backend;
pub use support::BackendSupport;
pub(crate) use support::CellException;
pub use support::DtypeBackendSupport;
pub(crate) use support::DtypeSupport;
pub use support::Support;
use support::dtype_serde;

use crate::pooling::Pooling;

/// The encoder architecture a model runs on.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum Family {
  Bert,
  Gemma3,
  GteNew,
  JinaBert,
  // Automatic snake_case would split this into "modern_bert"; the HF `model_type` and the page name are one word.
  #[serde(rename = "modernbert")]
  ModernBert,
  Mpnet,
  NomicBert,
  Qwen3,
  XlmRoberta,
  /// The flash XLM-RoBERTa encoder of `jinaai/jina-colbert-v2` and `jinaai/jina-reranker-v2-base-multilingual`; its
  /// checkpoints carry no `model_type`, so only a registry entry reaches it.
  XlmRobertaFlash,
}

impl Family {
  /// Maps an architecture key to a family: a `config.json` `model_type`, or `jina_bert` for a `bert` checkpoint
  /// with ALiBi positions; `XlmRobertaFlash` has no key.
  pub fn from_model_type(model_type: &str) -> Option<Family> {
    match model_type {
      "bert" => Some(Family::Bert),
      "gemma3_text" => Some(Family::Gemma3),
      "jina_bert" => Some(Family::JinaBert),
      "modernbert" => Some(Family::ModernBert),
      "mpnet" => Some(Family::Mpnet),
      "new" => Some(Family::GteNew),
      "nomic_bert" => Some(Family::NomicBert),
      "qwen3" => Some(Family::Qwen3),
      "xlm-roberta" => Some(Family::XlmRoberta),
      _ => None,
    }
  }
}

/// How a repository publishes the weights candding reads.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum WeightsFormat {
  /// `*.safetensors` files, memory-mapped.
  Safetensors,
  /// A PyTorch `pytorch_model.bin`, for `BAAI/bge-m3` and `prithivida/Splade_PP_en_v1`.
  Pickle,
  /// A raw NumPy array, for the head files of `Qdrant/minicoil-v1`.
  Npy,
  /// No weight file at all, for `Qdrant/bm25`.
  NoWeights,
}

/// The weight format of an id any registry carries, or of a sparse head's second repository; safetensors otherwise.
pub fn weights_format(id: &str) -> WeightsFormat {
  registered_weights_format(id).unwrap_or(WeightsFormat::Safetensors)
}

/// The dense, sparse, multi-vector and rerank registries in order, then a sparse head's second repository
/// (`Qdrant/minicoil-v1`), so fetching that id directly picks the right include patterns.
fn registered_weights_format(id: &str) -> Option<WeightsFormat> {
  entry_for(id)
    .map(|entry| entry.weights)
    .or_else(|| sparse::lookup(id).map(|d| d.weights))
    .or_else(|| multi_vector::lookup(id).map(|d| d.weights))
    .or_else(|| rerank::lookup(id).map(|d| d.weights))
    .or_else(|| sparse::head_source_format(id))
}

/// The registry's own spelling of an id any registry carries, matched case-insensitively; the cache is case-sensitive,
/// so an id is replaced with it before anything resolves or downloads.
pub fn canonical_id(id: &str) -> Option<&'static str> {
  entry_for(id)
    .map(|entry| entry.id)
    .or_else(|| sparse::canonical_id(id))
    .or_else(|| multi_vector::canonical_id(id))
    .or_else(|| rerank::canonical_id(id))
}

fn entry_for(id: &str) -> Option<&'static Entry> {
  entries().find(|entry| entry.id.eq_ignore_ascii_case(id))
}

/// Everything the library needs to know about a model beyond its weights. Build one with `new`, then assign fields.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ModelDescriptor {
  pub id: String,
  pub family: Family,
  pub dim: usize,
  pub default_max_length: usize,
  pub pooling: Pooling,
  pub normalize: bool,
  pub query_template: Option<String>,
  pub passage_template: Option<String>,
  pub default_instruction: Option<String>,
  pub dense_dirs: Vec<String>,
  pub mrl: bool,
  pub layer_norm_after_pooling: bool,
  pub license: String,
  pub size_gb: f32,
  pub gated: bool,
  pub phase: u8,
  pub backends: BackendSupport,
  /// The dtype the builder loads when the caller sets none; BF16 on CPU falls back to F32.
  #[serde(with = "dtype_serde")]
  pub default_dtype: DType,
  /// Per-dtype support on each device, derived from the entry's statuses; every builder refuses from this row.
  pub dtype_backends: DtypeBackendSupport,
}

impl ModelDescriptor {
  /// The four facts with no safe default; the rest is what a snapshot outside the catalog starts from: 512 tokens,
  /// unit-length rows, no template or Dense module, F32, untested everywhere, and catalog fields marked unknown.
  pub fn new(id: impl Into<String>, family: Family, dim: usize, pooling: Pooling) -> Self {
    Self {
      id: id.into(),
      family,
      dim,
      default_max_length: 512,
      pooling,
      normalize: true,
      query_template: None,
      passage_template: None,
      default_instruction: None,
      dense_dirs: Vec::new(),
      mrl: false,
      layer_norm_after_pooling: false,
      license: "unknown".to_string(),
      size_gb: 0.0,
      gated: false,
      phase: u8::MAX,
      backends: BackendSupport::untested(),
      default_dtype: DType::F32,
      dtype_backends: DtypeBackendSupport::untested(),
    }
  }
}

fn entries() -> impl Iterator<Item = &'static Entry> {
  bert::ENTRIES
    .iter()
    .chain(qwen3::ENTRIES.iter())
    .chain(gemma3::ENTRIES.iter())
    .chain(xlm_roberta::ENTRIES.iter())
    .chain(mpnet::ENTRIES.iter())
    .chain(jina_bert::ENTRIES.iter())
    .chain(nomic_bert::ENTRIES.iter())
    .chain(modernbert::ENTRIES.iter())
    .chain(gte_new::ENTRIES.iter())
}

#[cfg(test)]
mod tests;