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

use serde::Deserialize;

use super::position::PositionScheme;
use crate::error::CanddingError;
use crate::error::Result;
use crate::layers::activation::Activation;

/// Defaults mirror upstream's `XLMRobertaFlashConfig`; `type_vocab_size`, `max_position_embeddings` and
/// `position_embedding_type` have none, since a guess there would silently change the model.
#[derive(Clone, Debug, Deserialize)]
pub struct XlmRobertaFlashConfig {
  pub vocab_size: usize,
  pub hidden_size: usize,
  pub num_hidden_layers: usize,
  pub num_attention_heads: usize,
  pub intermediate_size: usize,
  pub type_vocab_size: usize,
  pub max_position_embeddings: usize,
  #[serde(default = "default_eps")]
  pub layer_norm_eps: f64,
  /// `None` when absent: the two implementation repositories default it differently (`rotary`, `absolute`),
  /// so the load checks it against the registry entry's scheme instead.
  #[serde(default)]
  pub position_embedding_type: Option<String>,
  /// Position ids start after this id; both implementation repositories default it to 1.
  #[serde(default = "default_pad_token_id")]
  pub pad_token_id: usize,
  #[serde(default = "default_rotary_base")]
  pub rotary_emb_base: f64,
  #[serde(default = "default_act")]
  pub hidden_act: Activation,
  /// Non-null only for a sibling checkpoint that trains LoRA adapters; refused, since merged weights would load under
  /// the plain tensor names and silently return the base model's embeddings.
  #[serde(default)]
  pub lora_adaptations: Option<serde_json::Value>,
  /// Refused: `true` expects a trained adapter on top of the base weights, which this family cannot apply.
  #[serde(default)]
  pub load_trained_adapters: bool,
  /// Refused: a Matryoshka checkpoint (jina-embeddings-v3) would load and silently return the base width.
  #[serde(default)]
  pub matryoshka_dimensions: Option<serde_json::Value>,
}

fn default_eps() -> f64 {
  1e-5
}

fn default_pad_token_id() -> usize {
  1
}

fn default_rotary_base() -> f64 {
  10_000.0
}

/// Not an upstream default: every published checkpoint declares the erf GELU.
fn default_act() -> Activation {
  Activation::Gelu
}

impl XlmRobertaFlashConfig {
  pub fn head_dim(&self) -> usize {
    self.hidden_size / self.num_attention_heads
  }

  /// The first position id of a learned table, as `XlmRobertaConfig::position_offset` numbers it.
  pub fn position_offset(&self) -> u32 {
    self.pad_token_id as u32 + 1
  }

  /// The longest sequence a learned table can index: 1026 rows carry 1024 positions.
  pub fn sequence_limit(&self) -> usize {
    self
      .max_position_embeddings
      .saturating_sub(self.pad_token_id + 1)
  }

  /// Refuses a `position_embedding_type` other than `rotary` or `absolute`, and a config that names
  /// `lora_adaptations`, `load_trained_adapters` or `matryoshka_dimensions`, which this struct would otherwise ignore.
  pub fn ensure_supported(&self) -> Result<()> {
    if let Some(value) = self.position_embedding_type.as_deref()
      && PositionScheme::from_config(value).is_none()
    {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "xlm-roberta-flash with {value} position embeddings"
      )));
    }
    if self.lora_adaptations.is_some() || self.load_trained_adapters {
      return Err(CanddingError::UnsupportedArchitecture(
        "xlm-roberta-flash with a LoRA adaptation declared (lora_adaptations/load_trained_adapters)"
          .to_string(),
      ));
    }
    if self.matryoshka_dimensions.is_some() {
      return Err(CanddingError::UnsupportedArchitecture(
        "xlm-roberta-flash with matryoshka_dimensions declared".to_string(),
      ));
    }
    Ok(())
  }
}

#[cfg(test)]
mod tests;