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 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use super::*;

/// `prithivida/Splade_PP_en_v1`'s real `modules.json` at `762be6a7206e2f299182705972a65e5c46e62be2`.
const SPLADE_MODULES: &str = r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.sparse_encoder.models.MLMTransformer"},{"idx":1,"name":"1","path":"1_SpladePooling","type":"sentence_transformers.sparse_encoder.models.SpladePooling"}]"#;

/// `BAAI/bge-m3`'s real `modules.json` at `5617a9f61b028005a4858fdac845db406aefb181`: no `Dense` module, nothing sparse.
const BGE_M3_MODULES: &str = 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"}]"#;

#[test]
fn splade_modules_are_accepted_by_the_sparse_pipeline() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"bert"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "model.safetensors", "");
  write(dir.path(), "modules.json", SPLADE_MODULES);
  let files = discover(dir.path(), Pipeline::Sparse).unwrap();
  assert_eq!(
    files.dense_dirs,
    Vec::<PathBuf>::new(),
    "SPLADE publishes no Dense module"
  );
}

#[test]
fn splade_modules_are_refused_by_the_dense_pipeline_naming_the_module() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"bert"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "model.safetensors", "");
  write(dir.path(), "modules.json", SPLADE_MODULES);
  let err = discover(dir.path(), Pipeline::Dense).unwrap_err();
  assert!(
    matches!(
      err,
      CanddingError::SparseOnlyModule(ref m)
        if m == "sentence_transformers.sparse_encoder.models.MLMTransformer"
    ),
    "{err}"
  );
  let message = err.to_string();
  assert!(
    message.contains("sentence_transformers.sparse_encoder.models.MLMTransformer"),
    "{message}"
  );
  assert!(message.contains("sparse pipeline"), "{message}");
}

#[test]
fn an_unrecognized_module_is_refused_by_both_pipelines() {
  let json = r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.models.Transformer"},{"idx":1,"name":"1","path":"1_LSTM","type":"sentence_transformers.models.LSTM"},{"idx":2,"name":"2","path":"2_Pooling","type":"sentence_transformers.models.Pooling"}]"#;
  for pipeline in [Pipeline::Dense, Pipeline::Sparse] {
    let dir = tempdir().unwrap();
    write(dir.path(), "config.json", r#"{"model_type":"bert"}"#);
    write(dir.path(), "tokenizer.json", "{}");
    write(dir.path(), "model.safetensors", "");
    write(dir.path(), "modules.json", json);
    let err = discover(dir.path(), pipeline).unwrap_err();
    assert!(
      matches!(err, CanddingError::UnsupportedModule(ref m) if m == "sentence_transformers.models.LSTM"),
      "{pipeline:?}: {err}"
    );
  }
}

#[test]
fn a_bge_m3_shaped_module_list_is_accepted_by_both_pipelines() {
  for pipeline in [Pipeline::Dense, Pipeline::Sparse] {
    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(), "modules.json", BGE_M3_MODULES);
    let files = discover(dir.path(), pipeline).unwrap();
    assert_eq!(
      files.dense_dirs,
      Vec::<PathBuf>::new(),
      "{pipeline:?}: bge-m3 publishes no Dense module"
    );
  }
}

/// `Qwen/Qwen3-Reranker-0.6B`'s real `modules.json` at `e61197ed45024b0ed8a2d74b80b4d909f1255473`; the 4B and 8B match it.
const QWEN3_RERANKER_MODULES: &str = r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.base.modules.transformer.Transformer"},{"idx":1,"name":"1","path":"1_LogitScore","type":"sentence_transformers.cross_encoder.modules.logit_score.LogitScore"}]"#;

fn discover_with(modules: &str, pipeline: Pipeline) -> Result<ModelFiles> {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"qwen3"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "model.safetensors", "");
  write(dir.path(), "modules.json", modules);
  discover(dir.path(), pipeline)
}

#[test]
fn the_logit_score_module_is_accepted_by_the_rerank_pipeline_alone() {
  let files = discover_with(QWEN3_RERANKER_MODULES, Pipeline::Rerank).unwrap();
  assert!(files.dense_dirs.is_empty());
  for pipeline in [Pipeline::Dense, Pipeline::Sparse] {
    match discover_with(QWEN3_RERANKER_MODULES, pipeline) {
      Err(CanddingError::UnsupportedModule(module)) => assert!(module.ends_with(".LogitScore")),
      other => panic!("{pipeline:?}: {other:?}"),
    }
  }
}

#[test]
fn the_rerank_pipeline_refuses_the_dense_modules() {
  match discover_with(BGE_M3_MODULES, Pipeline::Rerank) {
    Err(CanddingError::UnsupportedModule(module)) => assert!(module.ends_with(".Pooling")),
    other => panic!("{other:?}"),
  }
}