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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

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

use super::*;
use crate::layers::activation::Activation;

fn tiny_config() -> BertConfig {
  BertConfig {
    vocab_size: 32,
    hidden_size: 16,
    num_hidden_layers: 2,
    num_attention_heads: 2,
    intermediate_size: 32,
    hidden_act: Activation::Gelu,
    max_position_embeddings: 16,
    type_vocab_size: 2,
    layer_norm_eps: 1e-12,
  }
}

/// Randomly initialised weights are enough for shape and masking invariants.
fn tiny_model() -> BertModel {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  BertModel::load(tiny_config(), vb).unwrap()
}

fn run(model: &BertModel, rows: &[Vec<u32>]) -> Vec<Vec<Vec<f32>>> {
  let seq = rows.iter().map(Vec::len).max().unwrap();
  let mut ids = Vec::new();
  let mut mask = Vec::new();
  for row in rows {
    ids.extend(
      row
        .iter()
        .copied()
        .chain(std::iter::repeat_n(0, seq - row.len())),
    );
    mask
      .extend(std::iter::repeat_n(1u32, row.len()).chain(std::iter::repeat_n(0, seq - row.len())));
  }
  let input_ids = Tensor::from_vec(ids, (rows.len(), seq), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(mask, (rows.len(), seq), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let lengths: Vec<usize> = rows.iter().map(Vec::len).collect();
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  model.forward(&input).unwrap().to_vec3::<f32>().unwrap()
}

fn max_diff(a: &Tensor, b: &Tensor) -> f32 {
  (a - b)
    .unwrap()
    .abs()
    .unwrap()
    .flatten_all()
    .unwrap()
    .max(0)
    .unwrap()
    .to_scalar::<f32>()
    .unwrap()
}

#[test]
fn output_has_hidden_size_per_token() {
  let model = tiny_model();
  let out = run(&model, &[vec![2, 5, 7, 3]]);
  assert_eq!((out.len(), out[0].len(), out[0][0].len()), (1, 4, 16));
  assert_eq!(model.hidden_size(), 16);
  assert_eq!(model.max_position_embeddings(), 16);
}

#[test]
fn padding_does_not_change_real_tokens() {
  let model = tiny_model();
  let alone = run(&model, &[vec![2, 9, 3]]);
  let batched = run(&model, &[vec![2, 5, 7, 8, 3], vec![2, 9, 3]]);
  for (pos, token) in alone[0].iter().enumerate() {
    for (a, b) in token.iter().zip(&batched[1][pos]) {
      assert!((a - b).abs() < 1e-5, "position {pos}: {a} vs {b}");
    }
  }
}

/// Distinct random weights per layer, so the first and last layers' attention weights genuinely differ.
#[test]
fn forward_with_attention_returns_the_last_layers_weights() {
  let model = tiny_model();
  let input_ids = Tensor::from_vec(vec![2u32, 5, 7, 3], (1, 4), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32; 4], (1, 4), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let lengths = [4usize];
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };

  let (hidden, weights) = model.forward_with_attention(&input).unwrap();

  let embedded = model
    .embeddings
    .forward(&input_ids, &token_type_ids)
    .unwrap();
  let mask = AttentionMask::padding(
    &attention_mask,
    &lengths,
    model.config.num_attention_heads,
    model.config.head_dim(),
    embedded.dtype(),
  )
  .unwrap();
  let (_, first_layer_weights) = model.layers[0]
    .forward_with_weights(&embedded, &mask)
    .unwrap();
  let after_first = model.layers[0].forward(&embedded, &mask).unwrap();
  let (expected_hidden, last_layer_weights) = model.layers[1]
    .forward_with_weights(&after_first, &mask)
    .unwrap();

  assert!(
    max_diff(&hidden, &expected_hidden) < 1e-6,
    "hidden states must match a hand-replayed forward pass"
  );
  assert!(
    max_diff(&weights, &last_layer_weights) < 1e-6,
    "returned weights must match the last layer's own"
  );
  assert!(
    max_diff(&weights, &first_layer_weights) > 1e-3,
    "the two layers' attention must genuinely differ for this test to distinguish last from first"
  );
}