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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! A synthetic flash XLM-RoBERTa reranker: a zero-layer encoder with a learned position table and the head,
//! laid out as jina-reranker-v2 publishes it, with each numbering's logit computed independently in f64.

use std::collections::HashMap;

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;

pub(crate) const HIDDEN: usize = 4;
/// Covers `tiny_tokenizer_double_separator.json`'s 16 ids, so the builder can score through it.
pub(crate) const VOCAB: usize = 16;
/// The table's rows: `PAD + 1 = 2` the numbering skips, 6 usable positions.
pub(crate) const ROWS: usize = 8;
pub(crate) const PAD: usize = 1;

/// `config.json` naming `position_embedding_type`, and no `model_type`.
pub(crate) fn config_json(position_embedding_type: &str) -> String {
  format!(
    r#"{{"vocab_size":{VOCAB},"hidden_size":{HIDDEN},"num_hidden_layers":0,"num_attention_heads":1,
    "intermediate_size":1,"max_position_embeddings":{ROWS},"type_vocab_size":1,"layer_norm_eps":1e-5,
    "pad_token_id":{PAD},"position_embedding_type":"{position_embedding_type}"}}"#
  )
}

fn word_row(id: usize) -> [f32; HIDDEN] {
  let x = id as f32;
  [0.1 * x, -0.05 * x, 0.02 * x * x, 0.3]
}

fn position_row(row: usize) -> [f32; HIDDEN] {
  let x = row as f32;
  [x, 0.1 * x * x, -0.5 * x, 1.0 / (x + 1.0)]
}

const TOKEN_TYPE_ROW: [f32; HIDDEN] = [0.3, -0.2, 0.1, 0.05];
const DENSE: [[f32; HIDDEN]; HIDDEN] = [
  [0.9, -0.3, 0.2, 0.1],
  [0.4, 0.8, -0.5, 0.3],
  [-0.2, 0.6, 0.7, -0.4],
  [0.1, -0.7, 0.3, 0.5],
];
const DENSE_BIAS: [f32; HIDDEN] = [0.05, -0.1, 0.2, 0.0];
const OUT_PROJ: [f32; HIDDEN] = [0.7, -0.4, 0.25, 1.1];
const OUT_PROJ_BIAS: f32 = 0.1;

fn table(rows: usize, row: impl Fn(usize) -> [f32; HIDDEN]) -> Tensor {
  let values: Vec<f32> = (0..rows).flat_map(row).collect();
  Tensor::from_vec(values, (rows, HIDDEN), &Device::Cpu).unwrap()
}

/// Every tensor the flash encoder and the classification head read, `emb_ln` the identity affine.
pub(crate) fn tensors() -> HashMap<String, Tensor> {
  let vector = |values: &[f32]| Tensor::new(values, &Device::Cpu).unwrap();
  let mut map = HashMap::new();
  map.insert(
    "roberta.embeddings.word_embeddings.weight".to_string(),
    table(VOCAB, word_row),
  );
  map.insert(
    "roberta.embeddings.position_embeddings.weight".to_string(),
    table(ROWS, position_row),
  );
  map.insert(
    "roberta.embeddings.token_type_embeddings.weight".to_string(),
    table(1, |_| TOKEN_TYPE_ROW),
  );
  map.insert(
    "roberta.emb_ln.weight".to_string(),
    Tensor::ones(HIDDEN, DType::F32, &Device::Cpu).unwrap(),
  );
  map.insert(
    "roberta.emb_ln.bias".to_string(),
    Tensor::zeros(HIDDEN, DType::F32, &Device::Cpu).unwrap(),
  );
  map.insert(
    "classifier.dense.weight".to_string(),
    table(HIDDEN, |row| DENSE[row]),
  );
  map.insert("classifier.dense.bias".to_string(), vector(&DENSE_BIAS));
  map.insert(
    "classifier.out_proj.weight".to_string(),
    table(1, |_| OUT_PROJ),
  );
  map.insert(
    "classifier.out_proj.bias".to_string(),
    vector(&[OUT_PROJ_BIAS]),
  );
  map
}

/// Which table row the first position reads: from `pad + 1` (row `pad` for a pad id), from 0, or none.
#[derive(Clone, Copy, Debug)]
pub(crate) enum Numbering {
  FromPadPlusOne,
  FromZero,
  NoTable,
}

/// The raw logit of any pair whose first token is `first_id`: `out_proj(tanh(dense(emb_ln(word + position +
/// token_type))))` at the first position.
pub(crate) fn expected_logit(first_id: u32, numbering: Numbering) -> f64 {
  let position = match numbering {
    Numbering::FromPadPlusOne if first_id as usize == PAD => position_row(PAD),
    Numbering::FromPadPlusOne => position_row(PAD + 1),
    Numbering::FromZero => position_row(0),
    Numbering::NoTable => [0.0; HIDDEN],
  };
  logit(first_id, position)
}

/// The raw logit at the first position when its token is `first_id` and it reads `position` from the table.
fn logit(first_id: u32, position: [f32; HIDDEN]) -> f64 {
  let word = word_row(first_id as usize);
  let summed: [f64; HIDDEN] =
    std::array::from_fn(|k| word[k] as f64 + position[k] as f64 + TOKEN_TYPE_ROW[k] as f64);
  let mean = summed.iter().sum::<f64>() / HIDDEN as f64;
  let variance = summed.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / HIDDEN as f64;
  let normalized = summed.map(|v| (v - mean) / (variance + 1e-5).sqrt());
  let projected: [f64; HIDDEN] = std::array::from_fn(|r| {
    let dot: f64 = (0..HIDDEN)
      .map(|k| DENSE[r][k] as f64 * normalized[k])
      .sum();
    (dot + DENSE_BIAS[r] as f64).tanh()
  });
  (0..HIDDEN)
    .map(|k| OUT_PROJ[k] as f64 * projected[k])
    .sum::<f64>()
    + OUT_PROJ_BIAS as f64
}

/// The fixture keeps the numberings far apart, so a score held to the reference fails for either mistake.
#[test]
fn the_three_numberings_score_apart() {
  for first_id in [0, 1, 2] {
    let right = expected_logit(first_id, Numbering::FromPadPlusOne);
    for wrong in [Numbering::FromZero, Numbering::NoTable] {
      let moved = (right - expected_logit(first_id, wrong)).abs();
      assert!(moved > 0.05, "first id {first_id}, {wrong:?}: {moved}");
    }
  }
  let by_length = logit(PAD as u32, position_row(PAD + 1));
  let moved = (expected_logit(PAD as u32, Numbering::FromPadPlusOne) - by_length).abs();
  assert!(
    moved > 0.05,
    "the pad id at row pad against row pad + 1: {moved}"
  );
}