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 super::*;

fn architecture_of(body: &str) -> String {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", body);
  read_architecture(&dir.path().join("config.json")).unwrap()
}

/// `JinaBertEmbeddings.__init__` builds a position table only when `position_embedding_type` is not `"alibi"`.
#[test]
fn an_alibi_bert_resolves_to_the_jina_bert_architecture() {
  assert_eq!(
    architecture_of(r#"{"model_type":"bert","position_embedding_type":"alibi"}"#),
    "jina_bert"
  );
}

#[test]
fn a_plain_bert_keeps_its_model_type() {
  assert_eq!(architecture_of(r#"{"model_type":"bert"}"#), "bert");
  assert_eq!(
    architecture_of(r#"{"model_type":"bert","position_embedding_type":"absolute"}"#),
    "bert"
  );
}

#[test]
fn other_families_keep_their_model_type() {
  assert_eq!(
    architecture_of(r#"{"model_type":"xlm-roberta"}"#),
    "xlm-roberta"
  );
  assert_eq!(architecture_of(r#"{"model_type":"mpnet"}"#), "mpnet");
  assert_eq!(
    architecture_of(r#"{"model_type":"gemma3_text"}"#),
    "gemma3_text"
  );
  assert_eq!(
    architecture_of(r#"{"model_type":"roformer","position_embedding_type":"alibi"}"#),
    "roformer"
  );
}

#[test]
fn pooling_variants() {
  let dir = tempdir().unwrap();
  write(
    dir.path(),
    "mean.json",
    r#"{"pooling_mode_mean_tokens":true}"#,
  );
  write(
    dir.path(),
    "last.json",
    r#"{"pooling_mode_lasttoken":true}"#,
  );
  write(dir.path(), "empty.json", "{}");
  assert_eq!(
    read_pooling(&dir.path().join("mean.json")).unwrap(),
    Pooling::Mean
  );
  assert_eq!(
    read_pooling(&dir.path().join("last.json")).unwrap(),
    Pooling::LastToken
  );
  assert_eq!(
    read_pooling(&dir.path().join("empty.json")).unwrap(),
    Pooling::Mean
  );
}

#[test]
fn discovers_dense_directories_in_module_order() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"gemma3_text"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "model.safetensors", "");
  write(
    dir.path(),
    "modules.json",
    r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.models.Transformer"},{"idx":1,"name":"1","path":"1_Pooling","type":"sentence_transformers.models.Pooling"},{"idx":2,"name":"2","path":"2_Dense","type":"sentence_transformers.models.Dense"},{"idx":3,"name":"3","path":"3_Dense","type":"sentence_transformers.models.Dense"},{"idx":4,"name":"4","path":"4_Normalize","type":"sentence_transformers.models.Normalize"}]"#,
  );
  for name in ["2_Dense", "3_Dense"] {
    write(dir.path(), &format!("{name}/config.json"), "{}");
    write(dir.path(), &format!("{name}/model.safetensors"), "");
  }
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(
    files.dense_dirs,
    vec![dir.path().join("2_Dense"), dir.path().join("3_Dense")]
  );
  fs::remove_file(dir.path().join("3_Dense/model.safetensors")).unwrap();
  assert!(matches!(
    discover(dir.path(), Pipeline::Dense).unwrap_err(),
    CanddingError::MissingFile(p) if p.ends_with("3_Dense/model.safetensors")
  ));
}

#[test]
fn modules_json_reads_every_module_in_order() {
  let dir = tempdir().unwrap();
  write(
    dir.path(),
    "modules.json",
    r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.models.Transformer"},{"idx":1,"name":"1","path":"1_Pooling","type":"sentence_transformers.models.Pooling"},{"idx":2,"name":"2","path":"2_Normalize","type":"sentence_transformers.models.Normalize"}]"#,
  );
  let modules = read_modules(&dir.path().join("modules.json")).unwrap();
  assert_eq!(modules.len(), 3);
  assert_eq!(modules[1].path, "1_Pooling");
}

#[test]
fn sparse_linear_is_none_when_absent() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "pytorch_model.bin", "");
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(files.sparse_linear, None);
}

#[test]
fn sparse_linear_is_some_when_present() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "pytorch_model.bin", "");
  write(dir.path(), "sparse_linear.pt", "");
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(
    files.sparse_linear,
    Some(dir.path().join("sparse_linear.pt"))
  );
}

#[test]
fn colbert_linear_is_none_when_absent() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "pytorch_model.bin", "");
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(files.colbert_linear, None);
}

#[test]
fn colbert_linear_is_some_when_present_alongside_sparse_linear() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "pytorch_model.bin", "");
  write(dir.path(), "sparse_linear.pt", "");
  write(dir.path(), "colbert_linear.pt", "");
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(
    files.sparse_linear,
    Some(dir.path().join("sparse_linear.pt"))
  );
  assert_eq!(
    files.colbert_linear,
    Some(dir.path().join("colbert_linear.pt"))
  );
}