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
122
123

use candle_core::Device;
use candle_core::Tensor;

use super::*;

fn vocab(words: &[&str]) -> Vocab {
  let entries: Vec<String> = words.iter().map(|w| format!("{w:?}")).collect();
  let json = format!(
    r#"{{"vocab": [{}], "stem_mapping": {{}}}}"#,
    entries.join(", ")
  );
  Vocab::parse(&json).unwrap()
}

fn stopwords(words: &[&str]) -> Stopwords {
  Stopwords::parse(&words.join("\n"))
}

/// `tokens[id]` is that id's token text; real callers build it from the encoder's published tokenizer.
fn token_text(tokens: &[&str]) -> Vec<String> {
  tokens.iter().map(|t| t.to_string()).collect()
}

/// Its negative component survives: `SparseEmbedding::from_pairs` drops only exact zeros.
#[test]
fn a_single_recognized_word_end_to_end() {
  let vocab = vocab(&["cat"]);
  let stopwords = stopwords(&[]);
  let tokens = token_text(&["cat"]);
  let matrix = Tensor::from_vec(
    vec![
      0f32, 0., 0., 0., 0., 0., 0., 0., 1.0, -1.0, 0.5, 0.3, 0., 0., 0., 0.,
    ],
    (2, 2, 4),
    &Device::Cpu,
  )
  .unwrap();
  let hidden_row = vec![vec![1.0f32, 0.0]];

  let embedding = pooling::pool_row(
    &hidden_row,
    &[0],
    &tokens,
    &vocab,
    &stopwords,
    &matrix,
    false,
  )
  .unwrap();

  let t0 = 1.0f32.tanh();
  let t1 = (-1.0f32).tanh();
  let t2 = 0.5f32.tanh();
  let t3 = 0.3f32.tanh();
  let norm = (t0 * t0 + t1 * t1 + t2 * t2 + t3 * t3).sqrt();
  let tf = 1.0 * (1.2 + 1.0) / (1.0 + 1.2 * (1.0 - 0.75 + 0.75 * 1.0 / 150.0));

  assert_eq!(
    embedding.indices,
    vec![
      index::dense_index(1, 0),
      index::dense_index(1, 1),
      index::dense_index(1, 2),
      index::dense_index(1, 3),
    ]
  );
  let expected = [
    t0 / norm * tf,
    t1 / norm * tf,
    t2 / norm * tf,
    t3 / norm * tf,
  ];
  for (got, want) in embedding.values.iter().zip(expected) {
    assert!((got - want).abs() < 1e-5, "{got} vs {want}");
  }
}

/// The projection `[2, 0, 0, 1]` keeps two unequal components, so a dropped `tanh` changes their ratio.
#[test]
fn two_occurrences_of_the_same_word_exercise_averaging_and_term_frequency() {
  let vocab = vocab(&["cat"]);
  let stopwords = stopwords(&[]);
  let tokens = token_text(&["cat"]);
  let matrix = Tensor::from_vec(
    vec![
      0f32, 0., 0., 0., 0., 0., 0., 0., 3.0, 1.0, -1.0, 2.0, 1.0, -1.0, 1.0, 0.0,
    ],
    (2, 2, 4),
    &Device::Cpu,
  )
  .unwrap();
  let hidden_row = vec![vec![1.0f32, 0.0], vec![0.0, 1.0]];

  let embedding = pooling::pool_row(
    &hidden_row,
    &[0, 0],
    &tokens,
    &vocab,
    &stopwords,
    &matrix,
    false,
  )
  .unwrap();

  let t0 = 2.0f32.tanh();
  let t3 = 1.0f32.tanh();
  let norm = (t0 * t0 + t3 * t3).sqrt();
  let tf = 2.0 * (1.2 + 1.0) / (2.0 + 1.2 * (1.0 - 0.75 + 0.75 * 2.0 / 150.0));

  assert_eq!(
    embedding.indices,
    vec![index::dense_index(1, 0), index::dense_index(1, 3)]
  );
  let expected = [t0 / norm * tf, t3 / norm * tf];
  for (got, want) in embedding.values.iter().zip(expected) {
    assert!((got - want).abs() < 1e-5, "{got} vs {want}");
  }
}

mod query;
mod resplit;
mod rows;
mod shape;