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

//! ALiBi only acts inside an attention layer, so these tests need a real one.

use candle_nn::VarMap;

use super::*;
use crate::models::jina_bert::AlibiSlopeVariant;

/// Twelve heads, like `jinaai/jina-reranker-v1-tiny-en`: heads 8 to 11 take the tail the two variants disagree on.
const TWELVE_HEADS: &str = r#"{"vocab_size":32,"hidden_size":24,"num_hidden_layers":1,
  "num_attention_heads":12,"intermediate_size":16,"max_position_embeddings":64,"type_vocab_size":2,
  "layer_norm_eps":1e-12,"position_embedding_type":"alibi","feed_forward_type":"geglu"}"#;

/// `VarBuilder::get` zero-initializes the pooler and classifier, which would score every row zero.
fn fill_from_names(varmap: &VarMap) {
  for (name, var) in varmap.data().lock().unwrap().iter() {
    let seed = name.bytes().fold(2_166_136_261u32, |h, b| {
      (h ^ u32::from(b)).wrapping_mul(16_777_619)
    });
    let values: Vec<f32> = (0..var.elem_count() as u32)
      .map(|i| {
        let x = seed ^ i.wrapping_mul(2_654_435_761);
        (x >> 8) as f32 / (1u32 << 24) as f32 - 0.5
      })
      .collect();
    var
      .set(&Tensor::from_vec(values, var.shape(), &Device::Cpu).unwrap())
      .unwrap();
  }
}

fn twelve_head_score(varmap: &VarMap, alibi_variant: Option<AlibiSlopeVariant>) -> f32 {
  let vb = VarBuilder::from_varmap(varmap, DType::F32, &Device::Cpu);
  let model = load_rerank(
    "jina_bert",
    TWELVE_HEADS,
    vb,
    RerankConventions {
      head: RerankHead::BertPooler,
      score_activation: ScoreActivation::Identity,
      alibi_variant,
      position_scheme: None,
    },
  )
  .unwrap();
  let input_ids = Tensor::new(
    &[[1u32, 5, 9, 2, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 2]],
    &Device::Cpu,
  )
  .unwrap();
  let attention_mask = input_ids.ones_like().unwrap();
  let token_type_ids = Tensor::new(
    &[[0u32, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]],
    &Device::Cpu,
  )
  .unwrap();
  let lengths = vec![16];
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  model.forward(&input).unwrap()[0].value()
}

/// At twelve heads the halved tail scores this pair `4.2e-3` from the canonical slopes; `None` matches `Canonical`.
#[test]
fn the_alibi_variant_load_rerank_hands_a_twelve_head_encoder_changes_its_score() {
  let varmap = VarMap::new();
  twelve_head_score(&varmap, None);
  fill_from_names(&varmap);
  let canonical = twelve_head_score(&varmap, Some(AlibiSlopeVariant::Canonical));
  let halved = twelve_head_score(&varmap, Some(AlibiSlopeVariant::HalvedInterpolatedTail));
  let default = twelve_head_score(&varmap, None);
  assert!(
    (halved - canonical).abs() > 1e-4,
    "the halved tail must move the score: canonical {canonical}, halved {halved}"
  );
  assert_eq!(default, canonical, "None is the canonical slope formula");
}