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

//! Multi-head attention over one fused `Wqkv` projection, with rotary `q` and `k` and a bias-free output.

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

use super::config::NomicBertConfig;
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;

pub struct NomicBertAttention {
  q: Linear,
  k: Linear,
  v: Linear,
  out: Linear,
  heads: usize,
  head_dim: usize,
  scale: f64,
}

impl NomicBertAttention {
  /// `vb` is `encoder.layers.{i}`; `linear_no_bias` gives checkpoint-free unit tests a random init.
  pub fn load(config: &NomicBertConfig, vb: VarBuilder) -> Result<Self> {
    let hidden = config.n_embd;
    let fused = linear_no_bias(hidden, 3 * hidden, vb.pp("attn.Wqkv"))?;
    let (q, k, v) = split_fused_qkv(fused.weight(), hidden)?;
    Ok(Self {
      q: Linear::new(q, None),
      k: Linear::new(k, None),
      v: Linear::new(v, None),
      out: linear_no_bias(hidden, hidden, vb.pp("attn.out_proj"))?,
      heads: config.n_head,
      head_dim: config.head_dim(),
      scale: 1.0 / (config.head_dim() as f64).sqrt(),
    })
  }

  pub fn forward(
    &self,
    hidden: &Tensor,
    mask: &AttentionMask,
    rope: &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) = rope.apply(&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)?)
  }
}