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

use candle_nn::VarBuilder;

/// Present in `base-es` and `base-code`.
const QK_MARKER: &str = "encoder.layer.0.attention.self.layer_norm_q.weight";
/// Present in `base-code` alone.
const SPLIT_MARKER: &str = "encoder.layer.0.layer_norm_1.weight";

/// The checkpoint's block shape, probed from the safetensors header since `config.json` cannot tell them apart.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Variant {
  /// A LayerNorm over the full-width query and key projections, applied before the heads are split. The value
  /// projection is not normalized.
  pub qk_norm: bool,
  /// `layer_norm_1` and `layer_norm_2` in the layer, and an `up_gated_layer`/`down_layer` feed-forward with no
  /// norm and no residual of its own, whose activated half is the second rather than the first.
  pub split_norms: bool,
}

impl Variant {
  /// Layer 0's markers answer for the encoder; a wrong *present* fails the load, a wrong *absent* fails the goldens.
  pub fn probe(vb: &VarBuilder) -> Self {
    Self {
      qk_norm: vb.contains_tensor(QK_MARKER),
      split_norms: vb.contains_tensor(SPLIT_MARKER),
    }
  }
}

#[cfg(test)]
mod tests {
  use std::collections::HashMap;

  use candle_core::DType;
  use candle_core::Device;
  use candle_core::Tensor;

  use super::*;

  fn probe_with(names: &[&str]) -> Variant {
    let one = Tensor::ones(4, DType::F32, &Device::Cpu).unwrap();
    let map: HashMap<String, Tensor> = names
      .iter()
      .map(|n| ((*n).to_string(), one.clone()))
      .collect();
    Variant::probe(&VarBuilder::from_tensors(map, DType::F32, &Device::Cpu))
  }

  /// The configs differ only in `auto_map`, which candding must not follow, so the header is the only evidence.
  #[test]
  fn the_header_selects_the_block_shape() {
    assert_eq!(
      probe_with(&[]),
      Variant {
        qk_norm: false,
        split_norms: false
      },
      "small-en, base-en, base-de and base-zh"
    );
    assert_eq!(
      probe_with(&[QK_MARKER]),
      Variant {
        qk_norm: true,
        split_norms: false
      },
      "base-es"
    );
    assert_eq!(
      probe_with(&[QK_MARKER, SPLIT_MARKER]),
      Variant {
        qk_norm: true,
        split_norms: true
      },
      "base-code"
    );
    assert_eq!(
      probe_with(&[SPLIT_MARKER]),
      Variant {
        qk_norm: false,
        split_norms: true
      },
      "no published checkpoint, recorded so the probe's answer is defined rather than accidental"
    );
  }

  /// Every layer of every checkpoint carries these markers or none does, so layer 0 decides; a disagreement fails the load.
  #[test]
  fn the_markers_name_layer_zero() {
    assert_eq!(
      QK_MARKER,
      "encoder.layer.0.attention.self.layer_norm_q.weight"
    );
    assert_eq!(SPLIT_MARKER, "encoder.layer.0.layer_norm_1.weight");
    assert_eq!(
      probe_with(&[
        "encoder.layer.1.attention.self.layer_norm_q.weight",
        "encoder.layer.1.layer_norm_1.weight",
      ]),
      Variant {
        qk_norm: false,
        split_norms: false
      },
      "the probe reads layer 0 and nothing else"
    );
  }
}