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

use super::*;

/// Three features, since `LayerNorm` over two always normalizes to `±weight` however extreme the input.
#[test]
fn padded_positions_never_contribute() {
  let identity = Tensor::new(
    &[[1f32, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
    &Device::Cpu,
  )
  .unwrap();
  let zero_bias = Tensor::new(&[0f32, 0.0, 0.0], &Device::Cpu).unwrap();
  let norm_weight = Tensor::new(&[1f32, 1.0, 1.0], &Device::Cpu).unwrap();
  // Picks out the first of the three post-`LayerNorm` features as the (single-column) logit.
  let decoder_weight = Tensor::new(&[[1f32, 0.0, 0.0]], &Device::Cpu).unwrap();
  let decoder_bias = Tensor::new(&[0f32], &Device::Cpu).unwrap();

  let build = || {
    head(
      Linear::new(identity.clone(), Some(zero_bias.clone())),
      norm(norm_weight.clone(), zero_bias.clone(), 3, 1e-12),
      decoder_weight.clone(),
      decoder_bias.clone(),
      CHUNK,
    )
  };

  // The padded position is lopsided, so its first feature survives `LayerNorm` large enough to change the max.
  let alone_hidden = Tensor::new(&[[[0.3f32, 0.3, 0.3], [0.31, 0.29, 0.3]]], &Device::Cpu).unwrap();
  let alone_mask = Tensor::new(&[[1u32, 1]], &Device::Cpu).unwrap();
  let padded_hidden = Tensor::new(
    &[[[0.3f32, 0.3, 0.3], [0.31, 0.29, 0.3], [50.0, -50.0, -50.0]]],
    &Device::Cpu,
  )
  .unwrap();
  let padded_mask = Tensor::new(&[[1u32, 1, 0]], &Device::Cpu).unwrap();

  let alone = build().pool_logits(&alone_hidden, &alone_mask).unwrap();
  let padded = build().pool_logits(&padded_hidden, &padded_mask).unwrap();
  assert_rows_close(&alone, &padded);
}

/// Decoding without the transpose gives a different valid answer, so only the values catch it.
#[test]
fn the_decoder_is_the_embedding_transposed() {
  let transform_weight = Tensor::new(&[[0f32, 0.0], [0.0, 0.0]], &Device::Cpu).unwrap();
  let transform_bias = Tensor::new(&[0f32, 0.0], &Device::Cpu).unwrap();
  // Both features collapse to their LayerNorm bias, fixing the pre-decode vector at [1.0, 2.0].
  let transform_norm = norm(
    Tensor::new(&[1f32, 1.0], &Device::Cpu).unwrap(),
    Tensor::new(&[1f32, 2.0], &Device::Cpu).unwrap(),
    2,
    1e-12,
  );
  let decoder_weight = Tensor::new(&[[1f32, 10.0], [100.0, 1000.0]], &Device::Cpu).unwrap();
  let decoder_bias = Tensor::new(&[0f32, 0.0], &Device::Cpu).unwrap();
  let model = head(
    Linear::new(transform_weight, Some(transform_bias)),
    transform_norm,
    decoder_weight,
    decoder_bias,
    CHUNK,
  );

  let hidden = Tensor::new(&[[[0f32, 0.0]]], &Device::Cpu).unwrap();
  let mask = Tensor::new(&[[1u32]], &Device::Cpu).unwrap();
  let rows = model.pool_logits(&hidden, &mask).unwrap();

  // Transposed (correct): logits = [1,2] · row_j of decoder_weight = [1·1+2·10, 1·100+2·1000] = [21, 2100].
  // Untransposed would instead contract the other axis: [1·1+2·100, 1·10+2·1000] = [201, 2010].
  assert_eq!(rows[0].indices, vec![0, 1]);
  let expected = [22f32.ln(), 2101f32.ln()];
  for (got, want) in rows[0].values.iter().zip(expected) {
    assert!((got - want).abs() < 1e-4, "{got} vs {want}");
  }
}

#[test]
fn output_is_sorted_and_positive() {
  let transform = Linear::new(
    Tensor::new(&[[0f32]], &Device::Cpu).unwrap(),
    Some(Tensor::new(&[0f32], &Device::Cpu).unwrap()),
  );
  let transform_norm = norm(
    Tensor::new(&[1f32], &Device::Cpu).unwrap(),
    Tensor::new(&[2f32], &Device::Cpu).unwrap(),
    1,
    1e-12,
  );
  // logits = 2.0 · [-1, 1, 0.5] = [-2.0, 2.0, 1.0]; column 0's weight is exactly 0 after `relu` and drops out.
  let decoder_weight = Tensor::new(&[[-1f32], [1.0], [0.5]], &Device::Cpu).unwrap();
  let decoder_bias = Tensor::new(&[0f32, 0.0, 0.0], &Device::Cpu).unwrap();
  let model = head(
    transform,
    transform_norm,
    decoder_weight,
    decoder_bias,
    CHUNK,
  );

  let hidden = Tensor::new(&[[[0f32]]], &Device::Cpu).unwrap();
  let mask = Tensor::new(&[[1u32]], &Device::Cpu).unwrap();
  let rows = model.pool_logits(&hidden, &mask).unwrap();

  assert_eq!(rows[0].indices, vec![1, 2]);
  assert!(rows[0].indices.windows(2).all(|w| w[0] < w[1]));
  assert!(rows[0].values.iter().all(|&v| v > 0.0));
  let expected = [3f32.ln(), 2f32.ln()];
  for (got, want) in rows[0].values.iter().zip(expected) {
    assert!((got - want).abs() < 1e-5, "{got} vs {want}");
  }
}