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
127
128
129
130
131
132

use serde::Deserialize;

use crate::error::CanddingError;
use crate::error::Result;

/// The `NomicBertConfig` fields this family reads; the BERT-style aliases some configs also carry are ignored.
#[derive(Clone, Debug, Deserialize)]
pub struct NomicBertConfig {
  pub n_embd: usize,
  pub n_head: usize,
  pub n_layer: usize,
  pub n_inner: usize,
  pub activation_function: String,
  pub qkv_proj_bias: bool,
  pub mlp_fc1_bias: bool,
  pub mlp_fc2_bias: bool,
  pub prenorm: bool,
  pub causal: bool,
  pub rotary_emb_interleaved: bool,
  /// Absent from all three catalogued configs; the reference defaults it to `false`.
  #[serde(default)]
  pub norm_mlp: bool,
  pub type_vocab_size: usize,
  pub vocab_size: usize,
  /// Absent for the Snowflake checkpoint; `0` for the other two.
  #[serde(default)]
  pub pad_token_id: Option<u32>,
  pub layer_norm_epsilon: f64,
  pub rotary_emb_fraction: f64,
  /// No default: this family's 1000 is shared by no other rotary family, so a default would be a guess.
  pub rotary_emb_base: f64,
  /// A separate knob from `rotary_scaling_factor` below; all three catalogued checkpoints leave it `null`.
  #[serde(default)]
  pub rotary_emb_scale_base: Option<f64>,
  /// The dynamic-NTK factor: `2` for `v1` and Snowflake, `null` for `v1.5` (plain base at every length).
  #[serde(default)]
  pub rotary_scaling_factor: Option<f64>,
  /// The reference's default, which `v1` relies on by omitting the field.
  #[serde(default = "default_max_trained_positions")]
  pub max_trained_positions: usize,
  /// Absent for the Snowflake checkpoint.
  #[serde(default)]
  pub max_position_embeddings: Option<usize>,
}

fn default_max_trained_positions() -> usize {
  2048
}

impl NomicBertConfig {
  pub fn head_dim(&self) -> usize {
    self.n_embd / self.n_head
  }

  pub fn scaling_factor(&self) -> Option<f64> {
    self.rotary_scaling_factor
  }

  /// Refuses any config this port would load into the wrong tensors or arithmetic; all three checkpoints pass.
  pub fn ensure_supported(&self) -> Result<()> {
    if self.activation_function != "swiglu" {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: activation_function {:?} is not supported; this family reads swiglu checkpoints",
        self.activation_function
      )));
    }
    if self.prenorm {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: prenorm {} is not supported; the three catalogued checkpoints are post-norm, normalizing \
         after each residual add",
        self.prenorm
      )));
    }
    if self.causal {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: causal {} is not supported; the three catalogued checkpoints are bidirectional encoders",
        self.causal
      )));
    }
    if self.rotary_emb_fraction != 1.0 {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: rotary_emb_fraction {} is not supported; the three catalogued checkpoints rotate the full \
         head dimension",
        self.rotary_emb_fraction
      )));
    }
    if self.rotary_emb_interleaved {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: rotary_emb_interleaved {} is not supported; the three catalogued checkpoints rotate halves \
         rather than interleaved pairs",
        self.rotary_emb_interleaved
      )));
    }
    if self.qkv_proj_bias {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: qkv_proj_bias {} is not supported; the three catalogued checkpoints carry no bias on the \
         fused Wqkv projection",
        self.qkv_proj_bias
      )));
    }
    if self.mlp_fc1_bias {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: mlp_fc1_bias {} is not supported; the three catalogued checkpoints carry no bias on fc11 or \
         fc12",
        self.mlp_fc1_bias
      )));
    }
    if self.mlp_fc2_bias {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: mlp_fc2_bias {} is not supported; the down projection is loaded without a bias",
        self.mlp_fc2_bias
      )));
    }
    if let Some(scale_base) = self.rotary_emb_scale_base {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: rotary_emb_scale_base {scale_base} is not supported; the three catalogued checkpoints leave \
         this null"
      )));
    }
    if self.norm_mlp {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "nomic_bert: norm_mlp {} is not supported; the three catalogued checkpoints apply exactly one norm \
         around the mlp",
        self.norm_mlp
      )));
    }
    Ok(())
  }
}

#[cfg(test)]
pub(crate) mod tests;