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

use candle_core::Tensor;

use super::*;
use crate::models::nomic_bert::config::tests::V1;
use crate::models::nomic_bert::config::tests::V1_5;

/// The reference's `base * ((factor * seq / max_trained) - (factor - 1)) ^ (dim / (dim - 2))` at five lengths.
#[test]
fn the_dynamic_base_matches_the_reference_formula_at_the_tabulated_lengths() {
  for (seq, expected) in [
    (2048usize, 1000.00),
    (2049, 1001.01),
    (3000, 1971.04),
    (4096, 3108.22),
    (8192, 7453.48),
  ] {
    let base = dynamic_ntk_base(1000.0, 2.0, 2048, seq, 64);
    assert!(
      (base - expected).abs() < 0.01,
      "seq {seq}: {base} vs {expected}"
    );
  }
}

#[test]
fn the_base_is_unscaled_at_or_below_the_trained_length() {
  assert_eq!(dynamic_ntk_base(1000.0, 2.0, 2048, 2048, 64), 1000.0);
  assert_eq!(dynamic_ntk_base(1000.0, 2.0, 2048, 7, 64), 1000.0);
}

/// At 3000 tokens a `v1`-shaped table differs from the unscaled base-1000 one in the lowest frequency; `v1.5`'s equals it.
#[test]
fn only_a_configured_factor_scales_the_table() {
  let v1: NomicBertConfig = serde_json::from_str(V1).unwrap();
  let v1_5: NomicBertConfig = serde_json::from_str(V1_5).unwrap();
  let v1_rotary = NomicRotary::new(&v1, DType::F32, &Device::Cpu).unwrap();
  let v1_5_rotary = NomicRotary::new(&v1_5, DType::F32, &Device::Cpu).unwrap();
  let reference = RotaryEmbedding::new(64, 3000, 1000.0, DType::F32, &Device::Cpu).unwrap();

  let column = 64 / 2 - 1;
  let position = 2999;
  let value_at =
    |rope: &RotaryEmbedding| rope.table().0.to_vec2::<f32>().unwrap()[position][column];

  let reference_value = value_at(&reference);
  let v1_value = value_at(v1_rotary.for_length(3000).unwrap().as_rope());
  let v1_5_value = value_at(v1_5_rotary.for_length(3000).unwrap().as_rope());

  assert_eq!(
    v1_5_value, reference_value,
    "v1.5 col {column} pos {position}: {v1_5_value} vs unscaled reference {reference_value}"
  );
  assert!(
    (v1_value - reference_value).abs() > 0.1,
    "v1 col {column} pos {position}: {v1_value} vs unscaled reference {reference_value}, diff {}",
    (v1_value - reference_value).abs()
  );
}

#[test]
fn short_lengths_use_the_cached_table() {
  let v1: NomicBertConfig = serde_json::from_str(V1).unwrap();
  let v1_5: NomicBertConfig = serde_json::from_str(V1_5).unwrap();
  for config in [&v1, &v1_5] {
    let rotary = NomicRotary::new(config, DType::F32, &Device::Cpu).unwrap();
    assert!(matches!(rotary.for_length(100).unwrap(), Table::Cached(_)));
  }
}

#[test]
fn the_table_rotates_in_the_rotate_half_convention() {
  let v1_5: NomicBertConfig = serde_json::from_str(V1_5).unwrap();
  let rotary = NomicRotary::new(&v1_5, DType::F32, &Device::Cpu).unwrap();
  let table = rotary.for_length(5).unwrap();
  let rope = table.as_rope();

  let q = Tensor::randn(0f32, 1.0, (2, 3, 5, 64), &Device::Cpu).unwrap();
  let k = Tensor::randn(0f32, 1.0, (2, 1, 5, 64), &Device::Cpu).unwrap();
  let (rq, rk) = rope.apply(&q, &k).unwrap();
  let (cos, sin) = rope.table();
  let cos = cos.narrow(0, 0, 5).unwrap();
  let sin = sin.narrow(0, 0, 5).unwrap();
  let slow_q = candle_nn::rotary_emb::rope_slow(&q, &cos, &sin).unwrap();
  let slow_k = candle_nn::rotary_emb::rope_slow(&k, &cos, &sin).unwrap();

  let diff = |a: &Tensor, b: &Tensor| {
    (a - b)
      .unwrap()
      .abs()
      .unwrap()
      .flatten_all()
      .unwrap()
      .max(0)
      .unwrap()
      .to_scalar::<f32>()
      .unwrap()
  };
  assert!(diff(&rq, &slow_q) < 1e-5 && diff(&rk, &slow_k) < 1e-5);
  let first = |t: &Tensor| t.narrow(2, 0, 1).unwrap();
  assert!(
    diff(&first(&rq), &first(&q)) < 1e-6,
    "position 0 must be unchanged"
  );
}