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

//! GeGLU feed-forward over `layers::mlp::GatedMlp`: `Wi` splits into the activated half and the multiplier.

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

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

pub struct ModernBertMlp {
  gated: GatedMlp,
}

impl ModernBertMlp {
  /// `vb` is `layers.{i}`; `Wi`'s first `intermediate` rows are the gate, the rest the multiplier.
  pub fn load(config: &ModernBertConfig, vb: VarBuilder) -> Result<Self> {
    let hidden = config.hidden_size;
    let inter = config.intermediate_size;
    let wi = linear_no_bias(hidden, 2 * inter, vb.pp("mlp.Wi"))?;
    let activated = wi.weight().narrow(0, 0, inter)?.contiguous()?;
    let multiplier = wi.weight().narrow(0, inter, inter)?.contiguous()?;
    Ok(Self {
      gated: GatedMlp::new(
        Linear::new(activated, None),
        Linear::new(multiplier, None),
        linear_no_bias(inter, hidden, vb.pp("mlp.Wo"))?,
        Activation::Gelu,
      ),
    })
  }

  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 `ModernBertMlp` directly, so only `hidden_size` and `intermediate_size` 1 matter; the rest are gte-modernbert-base's.
  fn one_wide_config() -> ModernBertConfig {
    ModernBertConfig {
      hidden_size: 1,
      num_attention_heads: 1,
      num_hidden_layers: 1,
      intermediate_size: 1,
      hidden_activation: "gelu".to_string(),
      global_attn_every_n_layers: 3,
      local_attention: 128,
      global_rope_theta: 160_000.0,
      local_rope_theta: 10_000.0,
      attention_bias: false,
      mlp_bias: false,
      norm_bias: false,
      norm_eps: 1e-5,
      vocab_size: 50368,
      max_position_embeddings: 8192,
    }
  }

  /// The reference activates the FIRST half of `Wi`: `gelu_erf(1) * 3 = 2.524034`, where the second half
  /// would give `gelu_erf(3) * 1 = 2.995950` and tanh GELU `2.523576`.
  #[test]
  fn the_first_half_of_wi_is_activated() {
    let mut tensors = std::collections::HashMap::new();
    tensors.insert(
      "mlp.Wi.weight".to_string(),
      Tensor::new(&[[1f32], [3.0]], &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "mlp.Wo.weight".to_string(),
      Tensor::new(&[[1f32]], &Device::Cpu).unwrap(),
    );
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    let mlp = ModernBertMlp::load(&one_wide_config(), vb).unwrap();
    let out = mlp
      .forward(&Tensor::new(&[[1f32]], &Device::Cpu).unwrap())
      .unwrap();
    let value = out.to_vec2::<f32>().unwrap()[0][0];
    assert!((value - 2.524_034).abs() < 1e-5, "{value}");
  }
}