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

use candle_core::DType;
use candle_core::Device;
use candle_nn::VarMap;

use super::*;

const TINY_CONFIG: &str = r#"{"hidden_size":4,"vocab_size":8,"num_hidden_layers":0,
  "num_attention_heads":1,"intermediate_size":1,"hidden_act":"gelu",
  "max_position_embeddings":32,"type_vocab_size":1,"layer_norm_eps":1e-12}"#;

/// The `StanfordColbert` arm reads `tokenizer_path` first, so a `Tokenizer` error proves the dispatch reached it.
#[test]
fn stanford_colbert_reaches_its_real_arm_and_reads_the_tokenizer_it_is_given() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let missing_tokenizer = std::path::Path::new("/definitely/does/not/exist/tokenizer.json");
  let err = load_multi_vector(
    MultiVectorHead::StanfordColbert,
    TINY_CONFIG,
    vb,
    None,
    missing_tokenizer,
    MultiVectorConventions {
      keeps_cls_row: true,
      projection_in: 4,
      projection_out: 128,
      projection_bias: false,
      normalize: true,
      attends_to_mask_tokens: false,
      position_scheme: None,
    },
  )
  .err()
  .expect("StanfordColbert with a missing tokenizer path must still fail");
  match err {
    CanddingError::Tokenizer(_) => {}
    other => panic!("unexpected error for StanfordColbert: {other}"),
  }
}

/// The `M3Colbert` arm fails with `MissingFile` before parsing the config, which proves the dispatch reached it.
#[test]
fn m3_colbert_reaches_its_real_arm_and_asks_for_its_own_file() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let unused_tokenizer_path = std::path::Path::new("unused-for-m3-colbert.json");
  let err = load_multi_vector(
    MultiVectorHead::M3Colbert,
    TINY_CONFIG,
    vb,
    None,
    unused_tokenizer_path,
    MultiVectorConventions {
      keeps_cls_row: false,
      projection_in: 4,
      projection_out: 4,
      projection_bias: true,
      normalize: true,
      attends_to_mask_tokens: false,
      position_scheme: None,
    },
  )
  .err()
  .expect("M3Colbert with no colbert_linear path must still fail");
  match err {
    CanddingError::MissingFile(path) => {
      assert_eq!(path, std::path::PathBuf::from("colbert_linear.pt"));
    }
    other => panic!("unexpected error for M3Colbert: {other}"),
  }
}

#[test]
fn head_serializes_snake_case() {
  assert_eq!(
    serde_json::to_string(&MultiVectorHead::StanfordColbert).unwrap(),
    "\"stanford_colbert\""
  );
  assert_eq!(
    serde_json::from_str::<MultiVectorHead>("\"m3_colbert\"").unwrap(),
    MultiVectorHead::M3Colbert
  );
  assert_eq!(
    serde_json::to_string(&MultiVectorHead::JinaColbert).unwrap(),
    "\"jina_colbert\""
  );
  assert_eq!(
    serde_json::from_str::<MultiVectorHead>("\"jina_colbert\"").unwrap(),
    MultiVectorHead::JinaColbert
  );
}