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

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

use super::attention::XlmRobertaFlashAttention;
use super::config::XlmRobertaFlashConfig;
use crate::error::Result;
use crate::layers::activation::Activation;
use crate::layers::layer_norm::LayerNorm;
use crate::layers::mask::AttentionMask;
use crate::layers::rope::RotaryEmbedding;

/// A post-norm block, `LayerNorm(x + Sublayer(x))` after attention and after the feed-forward: the checkpoints build
/// `Block(prenorm=False)`, whatever the class default says, and carry no final norm a pre-norm stack would need.
pub struct XlmRobertaFlashLayer {
  attention: XlmRobertaFlashAttention,
  norm1: LayerNorm,
  fc1: Linear,
  activation: Activation,
  fc2: Linear,
  norm2: LayerNorm,
}

impl XlmRobertaFlashLayer {
  /// `vb` is one layer; the feed-forward is a plain biased `fc2(act(fc1(x)))` with no gate tensor.
  pub fn load(config: &XlmRobertaFlashConfig, vb: VarBuilder) -> Result<Self> {
    Ok(Self {
      attention: XlmRobertaFlashAttention::load(config, vb.pp("mixer"))?,
      norm1: LayerNorm::load(vb.pp("norm1"), config.hidden_size, config.layer_norm_eps)?,
      fc1: linear(
        config.hidden_size,
        config.intermediate_size,
        vb.pp("mlp.fc1"),
      )?,
      activation: config.hidden_act,
      fc2: linear(
        config.intermediate_size,
        config.hidden_size,
        vb.pp("mlp.fc2"),
      )?,
      norm2: LayerNorm::load(vb.pp("norm2"), config.hidden_size, config.layer_norm_eps)?,
    })
  }

  pub fn forward(
    &self,
    hidden: &Tensor,
    mask: &AttentionMask,
    rope: Option<&RotaryEmbedding>,
  ) -> Result<Tensor> {
    let attended = self.attention.forward(hidden, mask, rope)?;
    let hidden = self.norm1.forward(&(attended + hidden)?)?;
    let inner = self.activation.forward(&self.fc1.forward(&hidden)?)?;
    let mlp_out = self.fc2.forward(&inner)?;
    Ok(self.norm2.forward(&(mlp_out + hidden)?)?)
  }
}