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
113
114
115
116
117
118
119
120
121

use candle_core::Tensor;
use tempfile::TempDir;
use tempfile::tempdir;

use super::*;
use crate::models::EncoderInput;

const TRANSFORMER_POOLING: &str = r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.models.Transformer"},{"idx":1,"name":"1","path":"1_Pooling","type":"sentence_transformers.models.Pooling"}]"#;

const WITH_NORMALIZE: &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"}]"#;

/// A descriptor reads nothing from an encoder but its two sizes.
struct StubEncoder(usize);

impl Encoder for StubEncoder {
  fn forward(&self, _input: &EncoderInput) -> Result<Tensor> {
    unreachable!("a descriptor never runs the encoder")
  }

  fn hidden_size(&self) -> usize {
    384
  }

  fn max_position_embeddings(&self) -> usize {
    self.0
  }

  fn weights_dtype(&self) -> DType {
    DType::F32
  }
}

fn bare_snapshot() -> TempDir {
  let dir = tempdir().unwrap();
  fs::write(dir.path().join("config.json"), r#"{"model_type":"bert"}"#).unwrap();
  fs::write(dir.path().join("tokenizer.json"), "{}").unwrap();
  fs::write(dir.path().join("model.safetensors"), "").unwrap();
  dir
}

fn snapshot_descriptor(dir: &TempDir, encoder: &dyn Encoder) -> ModelDescriptor {
  let files = loader::discover(dir.path(), loader::Pipeline::Dense).unwrap();
  descriptor_from_snapshot("local", Family::Bert, &files, encoder, &[]).unwrap()
}

/// `sentence-transformers/all-MiniLM-L6-v2` publishes 256 in `sentence_bert_config.json` while its position table holds 512.
#[test]
fn a_snapshot_descriptor_takes_its_max_length_from_sentence_bert_config() {
  let dir = bare_snapshot();
  assert_eq!(
    snapshot_descriptor(&dir, &StubEncoder(512)).default_max_length,
    512,
    "without the file the rule is min(max_position_embeddings, 512)"
  );
  fs::write(
    dir.path().join("sentence_bert_config.json"),
    r#"{"max_seq_length": 256, "do_lower_case": false}"#,
  )
  .unwrap();
  assert_eq!(
    snapshot_descriptor(&dir, &StubEncoder(512)).default_max_length,
    256
  );
  fs::write(
    dir.path().join("sentence_bert_config.json"),
    r#"{"max_seq_length": 8192}"#,
  )
  .unwrap();
  assert_eq!(
    snapshot_descriptor(&dir, &StubEncoder(512)).default_max_length,
    512,
    "the position table still caps it"
  );
}

#[test]
fn a_snapshot_descriptor_normalizes_whatever_modules_json_lists() {
  let dir = bare_snapshot();
  assert!(
    snapshot_descriptor(&dir, &StubEncoder(512)).normalize,
    "no modules.json"
  );
  fs::write(dir.path().join("modules.json"), TRANSFORMER_POOLING).unwrap();
  assert!(
    snapshot_descriptor(&dir, &StubEncoder(512)).normalize,
    "a modules.json without a Normalize module"
  );
  fs::write(dir.path().join("modules.json"), WITH_NORMALIZE).unwrap();
  assert!(
    snapshot_descriptor(&dir, &StubEncoder(512)).normalize,
    "a modules.json with one"
  );
}

/// Every hidden state is the LayerNorm bias `[3, 4, 0, 0]`, so the pooled vector has norm 5.
#[test]
fn a_built_snapshot_returns_unit_rows_unless_normalize_is_turned_off() {
  let mut tensors = synthetic::zero_bert("", 16);
  let bias = Tensor::new(&[3f32, 4.0, 0.0, 0.0], &Device::Cpu).unwrap();
  tensors.insert("embeddings.LayerNorm.bias".to_string(), bias);
  let dir = synthetic::checkpoint(&synthetic::bert_config(16), "tiny_tokenizer.json", &tensors);
  fs::write(dir.path().join("modules.json"), TRANSFORMER_POOLING).unwrap();
  let spec = dir.path().to_str().unwrap();
  let row = |builder: TextEmbeddingBuilder| {
    let model = builder.device(Device::Cpu).build().unwrap();
    model.embed(&["hello world"], None).unwrap().remove(0)
  };
  for (builder, expected) in [
    (TextEmbedding::builder(spec), [0.6, 0.8, 0.0, 0.0]),
    (
      TextEmbedding::builder(spec).normalize(false),
      [3.0, 4.0, 0.0, 0.0],
    ),
  ] {
    let actual = row(builder);
    assert_eq!(actual.len(), expected.len());
    for (a, e) in actual.iter().zip(expected) {
      assert!((a - e).abs() < 1e-6, "{actual:?} against {expected:?}");
    }
  }
}