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::collections::HashMap;

use crate::registry;
use crate::registry::WeightsFormat;
use crate::registry::multi_vector;
use crate::registry::rerank;
use crate::registry::sparse;

fn every_entry() -> Vec<(&'static str, String, WeightsFormat)> {
  let mut entries: Vec<_> = registry::dense::all()
    .into_iter()
    .map(|d| {
      let weights = registry::entry_for(&d.id).expect("registered").weights;
      ("dense", d.id, weights)
    })
    .collect();
  entries.extend(
    sparse::all()
      .into_iter()
      .map(|d| ("sparse", d.id, d.weights)),
  );
  entries.extend(
    multi_vector::all()
      .into_iter()
      .map(|d| ("multi_vector", d.id, d.weights)),
  );
  entries.extend(
    rerank::all()
      .into_iter()
      .map(|d| ("rerank", d.id, d.weights)),
  );
  entries
}

#[test]
fn a_case_variant_of_a_rerank_only_id_canonicalizes_to_the_registry_spelling() {
  assert!(registry::dense::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").is_none());
  assert!(sparse::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").is_none());
  assert!(multi_vector::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").is_none());
  assert_eq!(
    registry::canonical_id("CROSS-ENCODER/ms-marco-minilm-l6-v2"),
    Some("cross-encoder/ms-marco-MiniLM-L6-v2")
  );
  for d in rerank::all() {
    assert_eq!(
      registry::canonical_id(&d.id.to_ascii_uppercase()),
      Some(d.id.as_str()),
      "{}",
      d.id
    );
  }
}

#[test]
fn a_case_variant_of_a_sparse_only_or_multi_vector_only_id_canonicalizes_too() {
  assert_eq!(
    registry::canonical_id("PRITHIVIDA/splade_pp_en_v1"),
    Some("prithivida/Splade_PP_en_v1")
  );
  assert_eq!(
    registry::canonical_id("COLBERT-IR/ColBERTv2.0"),
    Some("colbert-ir/colbertv2.0")
  );
  assert_eq!(registry::canonical_id("nope/nothing"), None);
}

/// Reads the lookup chain itself: `weights_format`'s safetensors fallback would answer the same for these ids.
#[test]
fn weights_format_returns_each_rerank_and_multi_vector_entrys_own_format() {
  for d in rerank::all() {
    assert_eq!(
      registry::registered_weights_format(&d.id),
      Some(d.weights),
      "{}",
      d.id
    );
    assert_eq!(registry::weights_format(&d.id), d.weights, "{}", d.id);
  }
  for d in multi_vector::all() {
    assert_eq!(
      registry::registered_weights_format(&d.id),
      Some(d.weights),
      "{}",
      d.id
    );
  }
  assert_eq!(registry::registered_weights_format("nope/nothing"), None);
}

/// Both stop at the first registry carrying an id, well defined only while every catalog spells and fetches it alike.
#[test]
fn an_id_in_several_registries_has_one_spelling_and_one_weights_format() {
  let mut seen: HashMap<String, Vec<(&'static str, String, WeightsFormat)>> = HashMap::new();
  for (catalog, id, weights) in every_entry() {
    seen
      .entry(id.to_ascii_lowercase())
      .or_default()
      .push((catalog, id, weights));
  }
  let shared: Vec<_> = seen.values().filter(|hits| hits.len() > 1).collect();
  assert!(
    !shared.is_empty(),
    "BAAI/bge-m3 is carried by three catalogs"
  );
  for hits in shared {
    let (_, first_id, first_weights) = &hits[0];
    for (catalog, id, weights) in hits {
      assert_eq!(id, first_id, "{catalog} spells {first_id} differently");
      assert_eq!(weights, first_weights, "{catalog} fetches {id} differently");
    }
  }
}