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

//! Word embeddings plus an optional one-row token-type term, then LayerNorm; positions are rotary.

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::GteNewConfig;
use crate::error::Result;
use crate::layers::layer_norm::LayerNorm;

/// `vb` is the model root, already past mGTE's `new.` prefix.
pub struct GteNewEmbeddings {
  word: Embedding,
  token_type: Option<Embedding>,
  norm: LayerNorm,
}

impl GteNewEmbeddings {
  /// `token_type` exists only when `type_vocab_size > 0`; en-v1.5 carries no such tensor.
  pub fn load(config: &GteNewConfig, vb: VarBuilder) -> Result<Self> {
    let embeddings = vb.pp("embeddings");
    let token_type = if config.has_token_types() {
      Some(embedding(
        config.type_vocab_size,
        config.hidden_size,
        embeddings.pp("token_type_embeddings"),
      )?)
    } else {
      None
    };
    Ok(Self {
      word: embedding(
        config.vocab_size,
        config.hidden_size,
        embeddings.pp("word_embeddings"),
      )?,
      token_type,
      norm: LayerNorm::load(
        embeddings.pp("LayerNorm"),
        config.hidden_size,
        config.layer_norm_eps,
      )?,
    })
  }

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

  /// Every position reads token-type row 0: the reference zeroes the ids when `type_vocab_size < 2` and larger
  /// tables are refused, so one broadcast row replaces the gather.
  pub fn forward(&self, input_ids: &Tensor) -> Result<Tensor> {
    let mut embeddings = self.word.forward(input_ids)?;
    if let Some(token_type) = &self.token_type {
      let row0 = token_type.embeddings().narrow(0, 0, 1)?;
      embeddings = embeddings.broadcast_add(&row0)?;
    }
    Ok(self.norm.forward(&embeddings)?)
  }
}

#[cfg(test)]
mod tests {
  use std::collections::HashMap;

  use candle_core::DType;
  use candle_core::Device;
  use candle_nn::VarBuilder;

  use super::*;

  fn config(type_vocab_size: usize) -> GteNewConfig {
    GteNewConfig {
      hidden_size: 4,
      num_hidden_layers: 1,
      num_attention_heads: 1,
      intermediate_size: 4,
      vocab_size: 10,
      type_vocab_size,
      hidden_act: "gelu".to_string(),
      layer_norm_type: "layer_norm".to_string(),
      layer_norm_eps: 1e-12,
      position_embedding_type: "rope".to_string(),
      rope_theta: 10_000.0,
      rope_scaling: None,
      max_position_embeddings: 32,
      pack_qkv: true,
      unpad_inputs: false,
      use_memory_efficient_attention: false,
    }
  }

  fn word_and_norm_tensors(hidden: usize, vocab: usize) -> HashMap<String, Tensor> {
    let mut tensors = HashMap::new();
    tensors.insert(
      "embeddings.word_embeddings.weight".to_string(),
      Tensor::zeros((vocab, hidden), DType::F32, &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "embeddings.LayerNorm.weight".to_string(),
      Tensor::ones(hidden, DType::F32, &Device::Cpu).unwrap(),
    );
    tensors.insert(
      "embeddings.LayerNorm.bias".to_string(),
      Tensor::zeros(hidden, DType::F32, &Device::Cpu).unwrap(),
    );
    tensors
  }

  #[test]
  fn a_zero_type_vocab_size_loads_without_a_token_type_tensor() {
    let tensors = word_and_norm_tensors(4, 10);
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    assert!(GteNewEmbeddings::load(&config(0), vb).is_ok());
  }

  #[test]
  fn a_positive_type_vocab_size_fails_to_load_without_a_token_type_tensor() {
    let tensors = word_and_norm_tensors(4, 10);
    let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
    assert!(GteNewEmbeddings::load(&config(1), vb).is_err());
  }
}