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

use super::*;

/// `head * 100 + i * 10 + j`, so a transposition or a head-order mistake shows in the numbers, not only the shape.
fn ramp_bias(heads: usize, seq: usize) -> Tensor {
  let mut values = Vec::with_capacity(heads * seq * seq);
  for h in 0..heads {
    for i in 0..seq {
      for j in 0..seq {
        values.push((h * 100 + i * 10 + j) as f32);
      }
    }
  }
  Tensor::from_vec(values, (1, heads, seq, seq), &Device::Cpu).unwrap()
}

fn floor_in(dtype: DType) -> f32 {
  Tensor::new(&[mask_min(dtype) as f32], &Device::Cpu)
    .unwrap()
    .to_dtype(dtype)
    .unwrap()
    .to_dtype(DType::F32)
    .unwrap()
    .to_vec1::<f32>()
    .unwrap()[0]
}

#[test]
fn padding_with_bias_keeps_the_bias_and_floors_the_padded_keys() {
  let attention_mask = Tensor::new(&[[1u32, 1, 0], [1, 1, 1]], &Device::Cpu).unwrap();
  let bias = ramp_bias(2, 3);
  let mask =
    AttentionMask::padding_with_bias(&attention_mask, &[2, 3], &bias, 2, 64, DType::F32).unwrap();
  let additive = mask
    .additive()
    .expect("padding_with_bias never returns none()");
  assert_eq!(additive.dims(), &[2, 2, 3, 3]);
  assert!(!mask.is_causal());
  let values = additive.flatten_all().unwrap().to_vec1::<f32>().unwrap();
  let at = |b: usize, h: usize, i: usize, j: usize| values[((b * 2 + h) * 3 + i) * 3 + j];
  let floor = floor_in(DType::F32);
  for h in 0..2 {
    for i in 0..3 {
      for j in 0..3 {
        let expected = (h * 100 + i * 10 + j) as f32;
        assert_eq!(at(1, h, i, j), expected, "row 1 ({h}, {i}, {j})");
        let expected0 = if j == 2 { floor } else { expected };
        assert_eq!(at(0, h, i, j), expected0, "row 0 ({h}, {i}, {j})");
      }
    }
  }
}

#[test]
fn padding_with_bias_keeps_one_row_when_nothing_is_padded() {
  let attention_mask = Tensor::new(&[[1u32, 1], [1, 1]], &Device::Cpu).unwrap();
  let bias = ramp_bias(3, 2);
  let mask =
    AttentionMask::padding_with_bias(&attention_mask, &[2, 2], &bias, 3, 64, DType::F32).unwrap();
  let additive = mask
    .additive()
    .expect("padding_with_bias never returns none()");
  assert_eq!(additive.dims(), &[1, 3, 2, 2]);
  assert_eq!(
    additive.flatten_all().unwrap().to_vec1::<f32>().unwrap(),
    bias.flatten_all().unwrap().to_vec1::<f32>().unwrap()
  );
  assert!(!mask.is_causal());
}

/// candle bails with `Mask type {mask_type:?} must match q type {itype:?}` otherwise.
#[test]
fn padding_with_bias_casts_once_at_the_end_in_every_dtype() {
  let attention_mask = Tensor::new(&[[1u32, 0]], &Device::Cpu).unwrap();
  let bias = ramp_bias(1, 2);
  for dtype in [DType::F32, DType::F16, DType::BF16] {
    let mask =
      AttentionMask::padding_with_bias(&attention_mask, &[1], &bias, 1, 64, dtype).unwrap();
    let additive = mask
      .additive()
      .expect("padding_with_bias never returns none()");
    assert_eq!(additive.dtype(), dtype, "{dtype:?}");
    let values = additive
      .to_dtype(DType::F32)
      .unwrap()
      .flatten_all()
      .unwrap()
      .to_vec1::<f32>()
      .unwrap();
    assert_eq!(values[0], 0.0, "{dtype:?} real key, bias 0");
    assert_eq!(values[1], floor_in(dtype), "{dtype:?} padded key");
    assert_eq!(values[2], 10.0, "{dtype:?} real key, bias 10");
    assert_eq!(values[3], floor_in(dtype), "{dtype:?} padded key");
  }
}

/// Past a magnitude of about sixteen, adding the bias to the F16 floor gives -inf, so a padded key must be selected.
#[test]
fn a_large_negative_bias_never_pushes_a_padded_key_past_the_floor() {
  let attention_mask = Tensor::new(&[[1u32, 0]], &Device::Cpu).unwrap();
  let bias = Tensor::from_vec(
    vec![-24.0f32, -24.0, -24.0, -24.0],
    (1, 1, 2, 2),
    &Device::Cpu,
  )
  .unwrap();
  for dtype in [DType::F16, DType::BF16] {
    let mask =
      AttentionMask::padding_with_bias(&attention_mask, &[1], &bias, 1, 64, dtype).unwrap();
    let values = mask
      .additive()
      .unwrap()
      .to_dtype(DType::F32)
      .unwrap()
      .flatten_all()
      .unwrap()
      .to_vec1::<f32>()
      .unwrap();
    assert!(
      values.iter().all(|v| v.is_finite()),
      "{dtype:?}: {values:?}"
    );
    assert_eq!(values[1], floor_in(dtype), "{dtype:?} padded key");
    assert!(
      (values[0] + 24.0).abs() < 0.2,
      "{dtype:?} real key kept the bias: {}",
      values[0]
    );
  }
}