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

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

use super::attention::Gemma3Attention;
use super::config::Gemma3Config;
use super::config::LayerType;
use crate::error::Result;
use crate::layers::mask::AttentionMask;
use crate::layers::mlp::GatedMlp;
use crate::layers::rms_norm::RmsNorm;
use crate::layers::rms_norm::RmsNormStyle;
use crate::layers::rope::RotaryEmbedding;

/// Per-type rotary tables (built at load) and masks (built per forward) the layers pick from.
pub struct LayerContext<'a> {
  pub rope_local: &'a RotaryEmbedding,
  pub rope_global: &'a RotaryEmbedding,
  pub mask_local: &'a AttentionMask,
  pub mask_global: &'a AttentionMask,
}

pub struct Gemma3Layer {
  layer_type: LayerType,
  input_layernorm: RmsNorm,
  self_attn: Gemma3Attention,
  post_attention_layernorm: RmsNorm,
  pre_feedforward_layernorm: RmsNorm,
  mlp: GatedMlp,
  post_feedforward_layernorm: RmsNorm,
}

impl Gemma3Layer {
  pub fn load(config: &Gemma3Config, layer_type: LayerType, vb: VarBuilder) -> Result<Self> {
    let norm = |name: &str| {
      RmsNorm::load(
        vb.pp(name),
        config.hidden_size,
        config.rms_norm_eps,
        RmsNormStyle::OnePlusScale,
      )
    };
    Ok(Self {
      layer_type,
      input_layernorm: norm("input_layernorm")?,
      self_attn: Gemma3Attention::load(config, vb.pp("self_attn"))?,
      post_attention_layernorm: norm("post_attention_layernorm")?,
      pre_feedforward_layernorm: norm("pre_feedforward_layernorm")?,
      mlp: GatedMlp::load(
        vb.pp("mlp"),
        config.hidden_size,
        config.intermediate_size,
        config.hidden_activation,
      )?,
      post_feedforward_layernorm: norm("post_feedforward_layernorm")?,
    })
  }

  /// Both sublayers are normalized before and after: `h += post(attn(pre(h)))`, then the same around the MLP.
  pub fn forward(&self, hidden: &Tensor, ctx: &LayerContext) -> Result<Tensor> {
    let (rope, mask) = match self.layer_type {
      LayerType::SlidingAttention => (ctx.rope_local, ctx.mask_local),
      LayerType::FullAttention => (ctx.rope_global, ctx.mask_global),
    };
    let attn = self
      .self_attn
      .forward(&self.input_layernorm.forward(hidden)?, rope, mask)?;
    let attended = (hidden + self.post_attention_layernorm.forward(&attn)?)?;
    let inner = self
      .mlp
      .forward(&self.pre_feedforward_layernorm.forward(&attended)?)?;
    Ok((attended + self.post_feedforward_layernorm.forward(&inner)?)?)
  }
}