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

//! Word embeddings and a norm only: positions come from rotary embeddings and token types are ignored.

use candle_core::DType;
use candle_core::Module;
use candle_core::Tensor;
use candle_nn::Embedding;
use candle_nn::VarBuilder;
use candle_nn::embedding;

use super::config::ModernBertConfig;
use crate::error::Result;
use crate::layers::layer_norm::LayerNorm;

pub struct ModernBertEmbeddings {
  tok: Embedding,
  norm: LayerNorm,
}

impl ModernBertEmbeddings {
  /// `vb` is the model's root `VarBuilder`.
  pub fn load(config: &ModernBertConfig, vb: VarBuilder) -> Result<Self> {
    let embeddings = vb.pp("embeddings");
    Ok(Self {
      tok: embedding(
        config.vocab_size,
        config.hidden_size,
        embeddings.pp("tok_embeddings"),
      )?,
      norm: LayerNorm::load_without_bias(
        embeddings.pp("norm"),
        config.hidden_size,
        config.norm_eps,
      )?,
    })
  }

  pub(super) fn weights_dtype(&self) -> DType {
    self.tok.embeddings().dtype()
  }

  pub fn forward(&self, input_ids: &Tensor) -> Result<Tensor> {
    Ok(self.norm.forward(&self.tok.forward(input_ids)?)?)
  }
}