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

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

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

mod fused_qkv;
mod position_scheme;
mod post_norm;
mod rotary;

fn tiny_config() -> XlmRobertaFlashConfig {
  XlmRobertaFlashConfig {
    vocab_size: 32,
    hidden_size: 8,
    num_hidden_layers: 2,
    num_attention_heads: 2,
    intermediate_size: 16,
    type_vocab_size: 1,
    max_position_embeddings: 32,
    layer_norm_eps: 1e-5,
    position_embedding_type: Some("rotary".to_string()),
    pad_token_id: 1,
    rotary_emb_base: 10_000.0,
    hidden_act: Activation::Gelu,
    lora_adaptations: None,
    load_trained_adapters: false,
    matryoshka_dimensions: None,
  }
}

/// `VarMap::set_one` writes into the storage the loaded module holds, so a test can rewrite one tensor in place.
fn tiny_model() -> (VarMap, XlmRobertaFlashModel) {
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = XlmRobertaFlashModel::load(tiny_config(), vb, Some(PositionScheme::Rotary)).unwrap();
  (varmap, model)
}

/// The absolute scheme numbers its table from the ids, so padding must use `tiny_config`'s `pad_token_id`.
fn run(model: &XlmRobertaFlashModel, rows: &[Vec<u32>]) -> Vec<Vec<Vec<f32>>> {
  let seq = rows.iter().map(Vec::len).max().unwrap();
  let pad = tiny_config().pad_token_id as u32;
  let mut ids = Vec::new();
  let mut mask = Vec::new();
  for row in rows {
    ids.extend(
      row
        .iter()
        .copied()
        .chain(std::iter::repeat_n(pad, 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_shift(a: &[Vec<Vec<f32>>], b: &[Vec<Vec<f32>>]) -> f32 {
  a.iter()
    .flatten()
    .flatten()
    .zip(b.iter().flatten().flatten())
    .map(|(x, y)| (x - y).abs())
    .fold(0.0, f32::max)
}

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

#[test]
fn padding_does_not_change_real_tokens() {
  let (_varmap, model) = tiny_model();
  let alone = run(&model, &[vec![3, 9, 5]]);
  let batched = run(&model, &[vec![3, 5, 7, 8, 9], vec![3, 9, 5]]);
  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}");
    }
  }
}

/// `type_vocab_size` is 1 here too, so a one-row token-type table must not be a shape error.
#[test]
fn a_one_row_token_type_table_loads_and_runs() {
  let config = tiny_config();
  assert_eq!(config.type_vocab_size, 1);
  let (_varmap, model) = tiny_model();
  let out = run(&model, &[vec![0, 4, 2]]);
  assert_eq!(out[0].len(), 3);
}