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

//! SwiGLU feed-forward over `layers::mlp::GatedMlp`: `fc2(fc11(x) * silu(fc12(x)))`, `fc12` the gate.

use candle_core::Tensor;
use candle_nn::VarBuilder;
use candle_nn::linear_no_bias;

use super::config::NomicBertConfig;
use crate::error::Result;
use crate::layers::activation::Activation;
use crate::layers::mlp::GatedMlp;

pub struct NomicBertMlp {
  gated: GatedMlp,
}

impl NomicBertMlp {
  /// `vb` is `encoder.layers.{i}`; `fc11`, `fc12` and `fc2` are separate bias-free tensors.
  pub fn load(config: &NomicBertConfig, vb: VarBuilder) -> Result<Self> {
    let hidden = config.n_embd;
    let intermediate = config.n_inner;
    let fc11 = linear_no_bias(hidden, intermediate, vb.pp("mlp.fc11"))?;
    let fc12 = linear_no_bias(hidden, intermediate, vb.pp("mlp.fc12"))?;
    let fc2 = linear_no_bias(intermediate, hidden, vb.pp("mlp.fc2"))?;
    Ok(Self {
      gated: GatedMlp::new(fc12, fc11, fc2, Activation::Silu),
    })
  }

  pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
    self.gated.forward(xs)
  }
}

#[cfg(test)]
mod tests {
  use candle_core::DType;
  use candle_core::Device;
  use candle_nn::VarBuilder;

  use super::*;

  /// Loads `NomicBertMlp` directly, so only `n_embd` and `n_inner` 1 matter.
  fn one_wide_config() -> NomicBertConfig {
    NomicBertConfig {
      n_embd: 1,
      n_head: 1,
      n_layer: 1,
      n_inner: 1,
      activation_function: "swiglu".to_string(),
      qkv_proj_bias: false,
      mlp_fc1_bias: false,
      mlp_fc2_bias: false,
      prenorm: false,
      causal: false,
      rotary_emb_interleaved: false,
      norm_mlp: false,
      type_vocab_size: 2,
      vocab_size: 8,
      pad_token_id: None,
      layer_norm_epsilon: 1e-12,
      rotary_emb_fraction: 1.0,
      rotary_emb_base: 1000.0,
      rotary_emb_scale_base: None,
      rotary_scaling_factor: None,
      max_trained_positions: 8,
      max_position_embeddings: None,
    }
  }

  /// `fc12` gates: `1 * silu(3) = 2.857722`, where gating `fc11` gives `3 * silu(1) = 2.193176`.
  #[test]
  fn the_block_gates_fc12() {
    let tensors = std::collections::HashMap::from([
      (
        "mlp.fc11.weight".to_string(),
        Tensor::new(&[[1f32]], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.fc12.weight".to_string(),
        Tensor::new(&[[3f32]], &Device::Cpu).unwrap(),
      ),
      (
        "mlp.fc2.weight".to_string(),
        Tensor::new(&[[1f32]], &Device::Cpu).unwrap(),
      ),
    ]);
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    let mlp = NomicBertMlp::load(&one_wide_config(), vb).unwrap();
    let hidden = Tensor::new(&[[1f32]], &Device::Cpu).unwrap();
    let value = mlp
      .forward(&hidden)
      .unwrap()
      .flatten_all()
      .unwrap()
      .to_vec1::<f32>()
      .unwrap()[0];
    assert!((value - 2.857_722).abs() < 1e-5, "{value}");
  }
}