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 std::borrow::Cow;

use candle_core::DType;

use crate::models::RerankHead;
use crate::models::Unembedding;
use crate::registry::CellException;
use crate::registry::DtypeSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::registry::rerank::entry::RerankEntry;
use crate::registry::t4::T4_F16_MISS_REFUSED;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PromptTemplate;

/// The default instruction: `prompts.query` of `config_sentence_transformers.json`, identical across the three sizes.
const DEFAULT_INSTRUCTION: &str =
  "Given a web search query, retrieve relevant passages that answer the query";

/// The Qwen3 rerankers' `LogitScore`: `true_token_id` 9693 ("yes"), `false_token_id` 2152 ("no"), from
/// `1_LogitScore/config.json`; the unembedding matrix differs by size.
const fn logit_score(unembedding: Unembedding) -> RerankHead {
  RerankHead::LogitScore {
    true_token_id: 9693,
    false_token_id: 2152,
    unembedding,
  }
}

/// One size's run results: `[cpu, metal, cuda]` statuses, dtype statuses, accelerator-only dtypes, exceptions.
struct Runs {
  devices: [Support; 3],
  dtypes: DtypeSupport,
  accel_only_dtypes: &'static [DType],
  cell_exceptions: &'static [CellException],
}

/// A Qwen3 reranker entry; sizes differ in id, size, unembedding, default dtype and runs.
const fn qwen3_reranker(
  id: &'static str,
  size_gb: f32,
  unembedding: Unembedding,
  default_dtype: DType,
  runs: Runs,
) -> RerankEntry {
  let [cpu, metal, cuda] = runs.devices;
  RerankEntry {
    id,
    family: Family::Qwen3,
    head: logit_score(unembedding),
    default_max_length: 8192,
    score_activation: ScoreActivation::Identity,
    input: RerankInput::Prompt {
      template: PromptTemplate::Qwen3Reranker,
      default_instruction: Cow::Borrowed(DEFAULT_INSTRUCTION),
    },
    alibi_variant: None,
    position_scheme: None,
    default_dtype,
    cpu,
    metal,
    cuda,
    dtypes: runs.dtypes,
    license: "Apache-2.0",
    size_gb,
    phase: 5,
    weights: WeightsFormat::Safetensors,
    gated: false,
    accel_only_dtypes: runs.accel_only_dtypes,
    cell_exceptions: runs.cell_exceptions,
  }
}

/// `Qwen/Qwen3-Reranker-0.6B`: `yes` minus `no` from the tied embedding table, the raw difference its `CrossEncoder`
/// reports; the 8192 default is the model card's. F16 misses its ceiling on the CPU and the T4.
pub(super) const QWEN3_RERANKER_0_6B: RerankEntry = qwen3_reranker(
  "Qwen/Qwen3-Reranker-0.6B",
  1.19,
  Unembedding::Tied,
  DType::F32,
  Runs {
    devices: [Support::Verified, Support::Verified, Support::Verified],
    dtypes: DtypeSupport::verified(),
    accel_only_dtypes: &[DType::F16],
    cell_exceptions: T4_F16_MISS_REFUSED,
  },
);

/// `Qwen/Qwen3-Reranker-4B`: the 0.6B's template, ids, instruction and activation over a larger Qwen3 with a tied
/// embedding table. BF16 by default and verified on Metal; F16 misses on both devices, and F32 has run on neither.
pub(super) const QWEN3_RERANKER_4B: RerankEntry = qwen3_reranker(
  "Qwen/Qwen3-Reranker-4B",
  8.04,
  Unembedding::Tied,
  DType::BF16,
  Runs {
    devices: [Support::Pending, Support::Verified, Support::Untested],
    dtypes: DtypeSupport {
      f32: Support::Untested,
      f16: Support::Unsupported,
      bf16: Support::Verified,
    },
    accel_only_dtypes: &[],
    cell_exceptions: &[],
  },
);

/// `Qwen/Qwen3-Reranker-8B`: logits from the untied `lm_head.weight`. BF16 by default, verified on Metal; F32 fits
/// neither device, and F16 is refused on the CPU and untested on Metal.
pub(super) const QWEN3_RERANKER_8B: RerankEntry = qwen3_reranker(
  "Qwen/Qwen3-Reranker-8B",
  16.38,
  Unembedding::Untied,
  DType::BF16,
  Runs {
    devices: [Support::Untested, Support::Verified, Support::Untested],
    dtypes: DtypeSupport {
      f32: Support::Untested,
      f16: Support::Untested,
      bf16: Support::Verified,
    },
    accel_only_dtypes: &[DType::F16],
    cell_exceptions: &[],
  },
);