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

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

use super::config::XlmRobertaFlashConfig;
use crate::error::Result;
use crate::layers::attention::scaled_dot_product;
use crate::layers::attention::split_fused_qkv;
use crate::layers::mask::AttentionMask;
use crate::layers::rope::RotaryEmbedding;

/// Multi-head attention over one fused, biased `Wqkv`, with rotary `q` and `k` under the rotary scheme.
pub struct XlmRobertaFlashAttention {
  q: Linear,
  k: Linear,
  v: Linear,
  out: Linear,
  heads: usize,
  head_dim: usize,
  scale: f64,
}

impl XlmRobertaFlashAttention {
  /// `vb` is one layer's `mixer`; the raw `Wqkv` is split by `split_fused_qkv` into query, key and value rows.
  /// Init hints match `candle_nn::linear`'s so `VarMap`-backed tests get a non-degenerate attention.
  pub fn load(config: &XlmRobertaFlashConfig, vb: VarBuilder) -> Result<Self> {
    let hidden = config.hidden_size;
    let wqkv = vb.pp("Wqkv");
    let bound = 1.0 / (hidden as f64).sqrt();
    let bias_init = candle_nn::Init::Uniform {
      lo: -bound,
      up: bound,
    };
    let fused_weight = wqkv.get_with_hints(
      (3 * hidden, hidden),
      "weight",
      candle_nn::init::DEFAULT_KAIMING_NORMAL,
    )?;
    let fused_bias = wqkv.get_with_hints(3 * hidden, "bias", bias_init)?;
    let (qw, kw, vw) = split_fused_qkv(&fused_weight, hidden)?;
    let (qb, kb, vbias) = split_fused_qkv(&fused_bias, hidden)?;
    Ok(Self {
      q: Linear::new(qw, Some(qb)),
      k: Linear::new(kw, Some(kb)),
      v: Linear::new(vw, Some(vbias)),
      out: linear(hidden, hidden, vb.pp("out_proj"))?,
      heads: config.num_attention_heads,
      head_dim: config.head_dim(),
      scale: 1.0 / (config.head_dim() as f64).sqrt(),
    })
  }

  /// Rotary positions run `0..seq` in every row, as in the reference; `rope` is `None` under the absolute scheme.
  pub fn forward(
    &self,
    hidden: &Tensor,
    mask: &AttentionMask,
    rope: Option<&RotaryEmbedding>,
  ) -> Result<Tensor> {
    let (batch, seq, _) = hidden.dims3()?;
    let split = |xs: Tensor| -> Result<Tensor> {
      Ok(
        xs.reshape((batch, seq, self.heads, self.head_dim))?
          .transpose(1, 2)?
          .contiguous()?,
      )
    };
    let q = split(self.q.forward(hidden)?)?;
    let k = split(self.k.forward(hidden)?)?;
    let v = split(self.v.forward(hidden)?)?;
    let (q, k) = match rope {
      Some(rope) => rope.apply(&q, &k)?,
      None => (q, k),
    };
    let context = scaled_dot_product(&q, &k, &v, mask, self.scale)?
      .transpose(1, 2)?
      .reshape((batch, seq, self.heads * self.head_dim))?;
    Ok(self.out.forward(&context)?)
  }
}