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

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

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

fn tiny_config() -> Qwen3Config {
  Qwen3Config {
    vocab_size: 32,
    hidden_size: 16,
    num_hidden_layers: 2,
    num_attention_heads: 4,
    num_key_value_heads: 2,
    head_dim: 4,
    intermediate_size: 32,
    hidden_act: Activation::Silu,
    max_position_embeddings: 16,
    rms_norm_eps: 1e-6,
    rope_theta: 10000.0,
    attention_bias: false,
    use_sliding_window: false,
    is_causal: true,
    tie_word_embeddings: false,
  }
}

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

fn run(model: &Qwen3Model, 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()
}

#[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 right_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}");
    }
  }
}

#[test]
fn later_tokens_do_not_change_earlier_ones() {
  let model = tiny_model();
  let a = run(&model, &[vec![2, 5, 7]]);
  let b = run(&model, &[vec![2, 5, 9]]);
  for pos in 0..2 {
    for (x, y) in a[0][pos].iter().zip(&b[0][pos]) {
      assert!((x - y).abs() < 1e-5, "position {pos} changed: {x} vs {y}");
    }
  }
  assert!(
    a[0][2]
      .iter()
      .zip(&b[0][2])
      .any(|(x, y)| (x - y).abs() > 1e-5),
    "the changed token must change its own state"
  );
}

#[test]
fn rejects_unsupported_variants() {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let mut gelu = tiny_config();
  gelu.hidden_act = Activation::Gelu;
  assert!(matches!(
    Qwen3Model::load(gelu, vb.clone()).unwrap_err(),
    CanddingError::UnsupportedArchitecture(_)
  ));
  let mut sliding = tiny_config();
  sliding.use_sliding_window = true;
  assert!(matches!(
    Qwen3Model::load(sliding, vb.clone()).unwrap_err(),
    CanddingError::UnsupportedArchitecture(_)
  ));
  let mut biased = tiny_config();
  biased.attention_bias = true;
  assert!(matches!(
    Qwen3Model::load(biased, vb).unwrap_err(),
    CanddingError::UnsupportedArchitecture(_)
  ));
}

mod bidirectional;