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
//! Free helpers for `SparseTextEmbeddingBuilder::build`: the MiniCOIL side repository, the weight files and the
//! refusal of a computed entry.

use std::path::PathBuf;

use candle_core::DType;

use crate::embed::fit;
use crate::error::CanddingError;
use crate::error::Result;
use crate::hub;
use crate::hub::ModelSource;
use crate::loader;
use crate::loader::MiniCoilSideFiles;
use crate::loader::ModelFiles;
use crate::models;
use crate::models::SparseHead;
use crate::registry::sparse::SparseComputation;
use crate::registry::sparse::SparseDescriptor;

/// Resolves `SparseHead::MiniCoil`'s side files from `head_source` at the pinned `head_revision` (`"main"` only
/// for a hand-built descriptor that leaves it unset); `None` for every other head.
pub(super) fn resolve_minicoil_side(
  descriptor: &SparseDescriptor,
) -> Result<Option<loader::MiniCoilSideFiles>> {
  if descriptor.head != models::SparseHead::MiniCoil {
    return Ok(None);
  }
  let head_source = descriptor
    .head_source
    .as_deref()
    .ok_or(CanddingError::MissingHeadSource)?;
  let revision = descriptor.head_revision.as_deref().unwrap_or("main");
  let source = ModelSource::parse(head_source, revision);
  let dir = hub::resolve(&source, crate::registry::WeightsFormat::Npy)?;
  Ok(Some(loader::discover_minicoil_files(&dir)?))
}

/// The files the build reads: the encoder's and bge-m3's `sparse_linear.pt` in the run's dtype, MiniCOIL's matrix in F32.
pub(super) fn weight_files(
  descriptor: &SparseDescriptor,
  files: &ModelFiles,
  minicoil: Option<&MiniCoilSideFiles>,
  dtype: DType,
) -> Vec<(PathBuf, DType)> {
  let mut weights = fit::at(&files.weights, dtype);
  if descriptor.head == SparseHead::M3Linear
    && let Some(head) = &files.sparse_linear
  {
    weights.push((head.clone(), dtype));
  }
  if let Some(side) = minicoil {
    weights.push((side.matrix.clone(), DType::F32));
  }
  weights
}

/// Refuses a `Computed` entry (`Qdrant/bm25`) before anything is fetched; `Bm25TextEmbedding::builder()` loads it.
pub(super) fn reject_computed_entry(descriptor: &SparseDescriptor) -> Result<()> {
  if descriptor.computation == SparseComputation::Computed {
    return Err(CanddingError::ComputedSparseModel {
      hint: "Bm25TextEmbedding::builder()".to_string(),
    });
  }
  Ok(())
}

#[cfg(test)]
mod tests;