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

use super::*;

#[test]
fn attends_to_mask_tokens_true_makes_every_position_attendable() {
  let model = identity_head(2, HashSet::new());
  assert!(model.attends_to_mask_tokens, "identity_head's own default");
  let attention_mask = Tensor::new(&[[1u32, 1, 1, 0, 0]], &Device::Cpu).unwrap();
  let effective = model.query_attention_mask(&attention_mask).unwrap();
  let values = effective.to_vec2::<u32>().unwrap();
  assert_eq!(
    values,
    vec![vec![1u32, 1, 1, 1, 1]],
    "attends_to_mask_tokens:true must make every position attendable, fill positions included"
  );
}

#[test]
fn attends_to_mask_tokens_false_leaves_the_mask_unchanged() {
  let model = head(
    Tensor::eye(2, DType::F32, &Device::Cpu).unwrap(),
    HashSet::new(),
    2,
    true,
    false,
  );
  assert!(!model.attends_to_mask_tokens);
  let attention_mask = Tensor::new(&[[1u32, 1, 1, 0, 0]], &Device::Cpu).unwrap();
  let effective = model.query_attention_mask(&attention_mask).unwrap();
  let values = effective.to_vec2::<u32>().unwrap();
  assert_eq!(
    values,
    vec![vec![1u32, 1, 1, 0, 0]],
    "attends_to_mask_tokens:false must leave the fill positions unattendable"
  );
}

#[test]
fn a_mask_with_no_padding_is_unaffected_by_either_setting() {
  let attention_mask = Tensor::new(&[[1u32, 1, 1, 1]], &Device::Cpu).unwrap();
  for attends in [true, false] {
    let model = head(
      Tensor::eye(2, DType::F32, &Device::Cpu).unwrap(),
      HashSet::new(),
      2,
      true,
      attends,
    );
    let effective = model.query_attention_mask(&attention_mask).unwrap();
    assert_eq!(
      effective.to_vec2::<u32>().unwrap(),
      vec![vec![1u32, 1, 1, 1]],
      "attends_to_mask_tokens:{attends} with nothing padded"
    );
  }
}

/// One real layer, so self-attention runs; the test compares two computations, so the random weights do not matter.
fn one_layer_encoder(hidden_size: usize) -> XlmRobertaFlashModel {
  let config = XlmRobertaFlashConfig {
    num_hidden_layers: 1,
    ..tiny_config(hidden_size)
  };
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  XlmRobertaFlashModel::load(config, vb, Some(PositionScheme::Rotary)).unwrap()
}

#[test]
fn forward_query_actually_calls_query_attention_mask() {
  // Eight wide, since a 2-wide LayerNorm output leaves the masked and unmasked cases an ULP apart.
  let model = JinaColbert {
    encoder: one_layer_encoder(8),
    weight: Tensor::eye(8, DType::F32, &Device::Cpu).unwrap(),
    bias: None,
    punctuation_ids: HashSet::new(),
    dim: 8,
    normalize: false,
    keeps_cls_row: true,
    attends_to_mask_tokens: true,
  };
  let input_ids = Tensor::from_vec(vec![5u32, 6, 0, 0], (1, 4), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32, 1, 0, 0], (1, 4), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let lengths = vec![2usize];
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };

  let actual = model.forward_query(&input).unwrap();

  let ones_mask = Tensor::ones_like(&attention_mask).unwrap();
  let all_attendable = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &ones_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  let expected_hidden = model.encoder.forward(&all_attendable).unwrap();
  let expected = model.select_query_rows(&expected_hidden).unwrap();

  let unchanged_hidden = model.encoder.forward(&input).unwrap();
  let unchanged = model.select_query_rows(&unchanged_hidden).unwrap();

  let actual_rows: Vec<f32> = actual[0].rows().flatten().copied().collect();
  let expected_rows: Vec<f32> = expected[0].rows().flatten().copied().collect();
  let unchanged_rows: Vec<f32> = unchanged[0].rows().flatten().copied().collect();

  for (a, e) in actual_rows.iter().zip(expected_rows.iter()) {
    assert!(
      (a - e).abs() < 1e-6,
      "forward_query must match an all-ones-mask forward pass bit for bit: {actual_rows:?} vs {expected_rows:?}"
    );
  }
  let differs_from_unchanged = actual_rows
    .iter()
    .zip(unchanged_rows.iter())
    .any(|(a, u)| (a - u).abs() > 1e-6);
  assert!(
    differs_from_unchanged,
    "forward_query must differ from a forward pass over the unchanged, still-padded mask: \
     {actual_rows:?} vs {unchanged_rows:?}"
  );
}