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

//! One pre-norm block with two residuals; layer 0 has no attention norm, as the embeddings' norm just ran.

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

use super::attention::ModernBertAttention;
use super::config::ModernBertConfig;
use super::mlp::ModernBertMlp;
use crate::error::Result;
use crate::layers::layer_norm::LayerNorm;
use crate::layers::mask::AttentionMask;
use crate::layers::rope::RotaryEmbedding;

pub struct ModernBertLayer {
  attn_norm: Option<LayerNorm>,
  attention: ModernBertAttention,
  mlp_norm: LayerNorm,
  mlp: ModernBertMlp,
}

impl ModernBertLayer {
  /// `vb` is `layers.{i}`.
  pub fn load(config: &ModernBertConfig, vb: VarBuilder, layer: usize) -> Result<Self> {
    let hidden = config.hidden_size;
    let attn_norm = if layer == 0 {
      None
    } else {
      Some(LayerNorm::load_without_bias(
        vb.pp("attn_norm"),
        hidden,
        config.norm_eps,
      )?)
    };
    Ok(Self {
      attn_norm,
      attention: ModernBertAttention::load(config, vb.clone())?,
      mlp_norm: LayerNorm::load_without_bias(vb.pp("mlp_norm"), hidden, config.norm_eps)?,
      mlp: ModernBertMlp::load(config, vb)?,
    })
  }

  pub fn forward(
    &self,
    hidden: &Tensor,
    mask: &AttentionMask,
    rope: &RotaryEmbedding,
  ) -> Result<Tensor> {
    let normed = match &self.attn_norm {
      Some(norm) => norm.forward(hidden)?,
      None => hidden.clone(),
    };
    let hidden = (hidden + self.attention.forward(&normed, mask, rope)?)?;
    let projected = self.mlp.forward(&self.mlp_norm.forward(&hidden)?)?;
    Ok((hidden + projected)?)
  }
}