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

//! Gated feed-forward over `GatedMlp`: the SECOND half of `up_gate_proj` is activated, unlike ModernBERT's `Wi`.

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

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

/// Rows `0..intermediate` of `up_gate_proj` are the raw multiplier, the rest the activated half.
pub struct GteNewMlp {
  gated: GatedMlp,
}

impl GteNewMlp {
  /// `vb` is `encoder.layer.{i}`; `up_gate_proj` has no bias, `down_proj` does.
  pub fn load(config: &GteNewConfig, vb: VarBuilder) -> Result<Self> {
    let hidden = config.hidden_size;
    let inter = config.intermediate_size;
    let up_gate = linear_no_bias(hidden, 2 * inter, vb.pp("mlp.up_gate_proj"))?;
    let up = up_gate.weight().narrow(0, 0, inter)?.contiguous()?;
    let gate = up_gate.weight().narrow(0, inter, inter)?.contiguous()?;
    Ok(Self {
      gated: GatedMlp::new(
        Linear::new(gate, None),
        Linear::new(up, None),
        linear(inter, hidden, vb.pp("mlp.down_proj"))?,
        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 `GteNewMlp` directly, so only `hidden_size` and `intermediate_size` 1 matter; the rest are gte-base-en-v1.5's.
  fn one_wide_config() -> GteNewConfig {
    GteNewConfig {
      hidden_size: 1,
      num_hidden_layers: 1,
      num_attention_heads: 1,
      intermediate_size: 1,
      vocab_size: 30_528,
      type_vocab_size: 0,
      hidden_act: "gelu".to_string(),
      layer_norm_type: "layer_norm".to_string(),
      layer_norm_eps: 1e-12,
      position_embedding_type: "rope".to_string(),
      rope_theta: 500_000.0,
      rope_scaling: None,
      max_position_embeddings: 8192,
      pack_qkv: true,
      unpad_inputs: false,
      use_memory_efficient_attention: false,
    }
  }

  /// The reference activates the SECOND half of `up_gate_proj`: `1 * gelu_erf(3) = 2.995950`, where the first half
  /// would give `3 * gelu_erf(1) = 2.524034` and tanh GELU `2.9963626`.
  #[test]
  fn the_second_half_of_up_gate_proj_is_activated() {
    let mut tensors = std::collections::HashMap::new();
    tensors.insert(
      "mlp.up_gate_proj.weight".to_string(),
      Tensor::new(&[[1f32], [3.0]], &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "mlp.down_proj.weight".to_string(),
      Tensor::new(&[[1f32]], &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "mlp.down_proj.bias".to_string(),
      Tensor::new(&[0f32], &Device::Cpu).unwrap(),
    );
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    let mlp = GteNewMlp::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.995_95).abs() < 1e-5, "{value}");
  }
}