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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

//! `prithivida/Splade_PP_en_v1`'s sparse head: a BERT MLM head max-pooled over the sequence.
//! `relu` and `log1p` apply to every logit before the max; the config's `activation_function` names only the `relu`.

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

use crate::error::Result;
use crate::layers::activation::Activation;
use crate::layers::layer_norm::LayerNorm;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::SparseModel;
use crate::models::bert::BertConfig;
use crate::models::bert::BertModel;
use crate::sparse::SparseEmbedding;

/// Positions per fold of the running maximum, so `(batch, seq, vocab)` never exists (476.9 MiB at batch 8).
const CHUNK: usize = 64;

/// `prithivida/Splade_PP_en_v1`'s sparse head, built on `bert::BertModel`.
pub struct SpladeSparse {
  bert: BertModel,
  transform: Linear,
  transform_norm: LayerNorm,
  /// `(vocab, hidden)`, tied to the word embeddings and shared via `BertModel::embed_tokens`; transposed to decode.
  decoder_weight: Tensor,
  /// `cls.predictions.bias`, tied to `cls.predictions.decoder.bias`; either name reads the same values.
  decoder_bias: Tensor,
  vocab_size: usize,
  chunk: usize,
}

impl SpladeSparse {
  /// Builds the BERT encoder at `bert.` and the MLM head's own tensors at `cls.predictions.`.
  pub fn load(config: BertConfig, vb: VarBuilder) -> Result<Self> {
    let hidden_size = config.hidden_size;
    let vocab_size = config.vocab_size;
    let eps = config.layer_norm_eps;
    let bert = BertModel::load(config, vb.pp("bert"))?;
    let transform_vb = vb.pp("cls").pp("predictions").pp("transform");
    let transform = linear(hidden_size, hidden_size, transform_vb.pp("dense"))?;
    let transform_norm = LayerNorm::load(transform_vb.pp("LayerNorm"), hidden_size, eps)?;
    let decoder_weight = bert.embed_tokens().clone();
    let decoder_bias = vb.pp("cls").pp("predictions").get(vocab_size, "bias")?;
    Ok(Self {
      bert,
      transform,
      transform_norm,
      decoder_weight,
      decoder_bias,
      vocab_size,
      chunk: CHUNK,
    })
  }

  /// The MLM head's logits, `transform_norm(gelu_erf(transform(h))) · decoder_weightᵀ + decoder_bias`.
  fn chunk_logits(&self, h_chunk: &Tensor) -> Result<Tensor> {
    let t = self.transform.forward(h_chunk)?;
    let t = Activation::Gelu.forward(&t)?;
    let t = self.transform_norm.forward(&t)?;
    let decoder = Linear::new(self.decoder_weight.clone(), Some(self.decoder_bias.clone()));
    Ok(decoder.forward(&t)?)
  }

  /// Folds a running `(batch, vocab)` maximum chunk by chunk. Padded logits are masked to 0 before `relu`,
  /// since a padded logit can have any sign.
  fn pool_logits(&self, hidden: &Tensor, attention_mask: &Tensor) -> Result<Vec<SparseEmbedding>> {
    let (_, seq, _) = hidden.dims3()?;
    let mut pooled: Option<Tensor> = None;
    let mut start = 0;
    while start < seq {
      let len = self.chunk.min(seq - start);
      let h_chunk = hidden.narrow(1, start, len)?;
      let mask_chunk = attention_mask
        .narrow(1, start, len)?
        .to_dtype(hidden.dtype())?
        .unsqueeze(2)?;
      let masked = self.chunk_logits(&h_chunk)?.broadcast_mul(&mask_chunk)?;
      // candle has no `log1p`; `(1 + x).ln()` is accurate enough after `relu`.
      let weight = masked.relu()?.affine(1.0, 1.0)?.log()?;
      let chunk_max = weight.max(1)?;
      pooled = Some(match pooled {
        Some(prev) => prev.broadcast_maximum(&chunk_max)?,
        None => chunk_max,
      });
      start += len;
    }
    let pooled =
      pooled.expect("hidden carries at least one position: every EncoderInput has a CLS token");
    let rows = pooled.to_dtype(DType::F32)?.to_vec2::<f32>()?;
    Ok(
      rows
        .into_iter()
        .map(|row| {
          SparseEmbedding::from_pairs(row.into_iter().enumerate().map(|(i, v)| (i as u32, v)))
        })
        .collect(),
    )
  }
}

impl SparseModel for SpladeSparse {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let hidden = self.bert.forward(input)?;
    self.pool_logits(&hidden, input.attention_mask)
  }

  fn vocab_size(&self) -> usize {
    self.vocab_size
  }

  fn max_position_embeddings(&self) -> usize {
    self.bert.max_position_embeddings()
  }

  fn weights_dtype(&self) -> DType {
    self.bert.weights_dtype()
  }
}

#[cfg(test)]
mod tests;