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

use super::*;

/// `Alibaba-NLP/gte-modernbert-base`'s published `config.json`, verbatim (1184 bytes).
const PUBLISHED: &str = r#"{
  "architectures": ["ModernBertModel"],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 50281,
  "classifier_activation": "gelu",
  "classifier_bias": false,
  "classifier_dropout": 0.0,
  "classifier_pooling": "mean",
  "cls_token_id": 50281,
  "decoder_bias": true,
  "deterministic_flash_attn": false,
  "embedding_dropout": 0.0,
  "eos_token_id": 50282,
  "global_attn_every_n_layers": 3,
  "global_rope_theta": 160000.0,
  "gradient_checkpointing": false,
  "hidden_activation": "gelu",
  "hidden_size": 768,
  "initializer_cutoff_factor": 2.0,
  "initializer_range": 0.02,
  "intermediate_size": 1152,
  "layer_norm_eps": 1e-05,
  "local_attention": 128,
  "local_rope_theta": 10000.0,
  "max_position_embeddings": 8192,
  "mlp_bias": false,
  "mlp_dropout": 0.0,
  "model_type": "modernbert",
  "norm_bias": false,
  "norm_eps": 1e-05,
  "num_attention_heads": 12,
  "num_hidden_layers": 22,
  "pad_token_id": 50283,
  "position_embedding_type": "absolute",
  "sep_token_id": 50282,
  "sparse_pred_ignore_index": -100,
  "sparse_prediction": false,
  "torch_dtype": "float16",
  "transformers_version": "4.48.0.dev0",
  "vocab_size": 50368
}"#;

fn published() -> ModernBertConfig {
  serde_json::from_str(PUBLISHED).unwrap()
}

#[test]
fn reads_the_published_config() {
  let config = published();
  assert_eq!(config.hidden_size, 768);
  assert_eq!(config.num_attention_heads, 12);
  assert_eq!(config.num_hidden_layers, 22);
  assert_eq!(config.intermediate_size, 1152);
  assert_eq!(config.hidden_activation, "gelu");
  assert_eq!(config.global_attn_every_n_layers, 3);
  assert_eq!(config.local_attention, 128);
  assert_eq!(config.global_rope_theta, 160000.0);
  assert_eq!(config.local_rope_theta, 10000.0);
  assert!(!config.attention_bias);
  assert!(!config.mlp_bias);
  assert!(!config.norm_bias);
  assert_eq!(config.norm_eps, 1e-5);
  assert_eq!(config.vocab_size, 50368);
  assert_eq!(config.max_position_embeddings, 8192);
  assert_eq!(config.head_dim(), 64);
  config.ensure_supported().unwrap();
}

/// Layer `i` is global exactly when `i % global_attn_every_n_layers == 0`: 0, 3, 6, … and the last layer 21.
#[test]
fn every_third_layer_is_global_starting_with_the_first() {
  let config = published();
  assert!(config.is_global(0));
  assert!(!config.is_global(1));
  assert!(!config.is_global(2));
  assert!(config.is_global(3));
  assert!(config.is_global(21));
  assert_eq!((0..22).filter(|&i| config.is_global(i)).count(), 8);
}

/// The reference keeps `|i - j| <= local_attention / 2`, `AttentionMask::sliding` keeps `|i - j| < window`,
/// so the window is `local_attention / 2 + 1`: 65 for the published 128.
#[test]
fn the_attention_window_is_half_the_span_plus_one() {
  assert_eq!(published().attention_window(), 65);
}

#[test]
fn each_refusal_names_the_field_and_the_value() {
  let cases = [
    (
      "\"hidden_activation\": \"gelu\"",
      "\"hidden_activation\": \"silu\"",
      "hidden_activation silu",
    ),
    (
      "\"attention_bias\": false",
      "\"attention_bias\": true",
      "attention_bias true",
    ),
    ("\"mlp_bias\": false", "\"mlp_bias\": true", "mlp_bias true"),
    (
      "\"norm_bias\": false",
      "\"norm_bias\": true",
      "norm_bias true",
    ),
    (
      "\"global_attn_every_n_layers\": 3",
      "\"global_attn_every_n_layers\": 0",
      "global_attn_every_n_layers 0",
    ),
    (
      "\"local_attention\": 128",
      "\"local_attention\": 127",
      "local_attention 127",
    ),
  ];
  for (from, to, needle) in cases {
    assert!(PUBLISHED.contains(from), "{from}");
    let config: ModernBertConfig = serde_json::from_str(&PUBLISHED.replace(from, to)).unwrap();
    let err = config.ensure_supported().unwrap_err().to_string();
    assert!(err.contains(needle), "{needle}: {err}");
  }
}