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

//! Unlike the bm25 and bm42 query paths, minicoil's runs the encoder.

use std::collections::HashMap;

use candle_nn::VarBuilder;

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

/// A zero-layer encoder with hand-set tensors, so the hidden state is exact: the specials' rows are zero and drop,
/// and "cat"'s `[1.0, 0.0]` normalizes to `[1.0, -1.0]` under population variance.
#[test]
fn the_query_path_weighs_uniformly_where_the_passage_path_does_not() {
  let config = JinaBertConfig {
    vocab_size: 3,
    hidden_size: 2,
    num_hidden_layers: 0,
    num_attention_heads: 1,
    intermediate_size: 2,
    type_vocab_size: 2,
    layer_norm_eps: 1e-12,
    max_position_embeddings: 16,
    position_embedding_type: "alibi".to_string(),
    feed_forward_type: "geglu".to_string(),
  };
  let tensors = HashMap::from([
    (
      "embeddings.word_embeddings.weight".to_string(),
      Tensor::from_vec(vec![0f32, 0., 1.0, 0., 0., 0.], (3, 2), &Device::Cpu).unwrap(),
    ),
    (
      "embeddings.token_type_embeddings.weight".to_string(),
      Tensor::from_vec(vec![0f32, 0., 0., 0.], (2, 2), &Device::Cpu).unwrap(),
    ),
    (
      "embeddings.LayerNorm.weight".to_string(),
      Tensor::from_vec(vec![1f32, 1.], 2, &Device::Cpu).unwrap(),
    ),
    (
      "embeddings.LayerNorm.bias".to_string(),
      Tensor::from_vec(vec![0f32, 0.], 2, &Device::Cpu).unwrap(),
    ),
  ]);
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let encoder = JinaBertModel::load(config, vb, AlibiSlopeVariant::Canonical).unwrap();

  // Both hidden dimensions add +2.0 to every component, so all four stay positive and equal (0.5) after normalization.
  let matrix = Tensor::from_vec(
    vec![
      0f32, 0., 0., 0., 0., 0., 0., 0., 1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0,
    ],
    (2, 2, 4),
    &Device::Cpu,
  )
  .unwrap();
  let model = MiniCoilSparse::new(
    Box::new(encoder),
    matrix,
    vocab(&["cat"]),
    stopwords(&[]),
    token_text(&["[CLS]", "cat", "[SEP]"]),
  )
  .unwrap();

  let input_ids = Tensor::from_vec(vec![0u32, 1, 2], (1, 3), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32, 1, 1], (1, 3), &Device::Cpu).unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &[3],
  };

  let passage = model.forward(&input).unwrap();
  let query = model.forward_query(&input).unwrap();

  let tf: f32 = 1.0 * (1.2 + 1.0) / (1.0 + 1.2 * (1.0 - 0.75 + 0.75 * 1.0 / 150.0));
  assert!(
    (tf - 1.0).abs() > 1e-3,
    "tf must not coincidentally be 1.0: {tf}"
  );

  let expected_indices: Vec<u32> = (0..4).map(|c| index::dense_index(1, c)).collect();
  assert_eq!(passage[0].indices, expected_indices);
  assert_eq!(query[0].indices, expected_indices);
  for (&p, &q) in passage[0].values.iter().zip(&query[0].values) {
    assert!(
      (q - 0.5).abs() < 1e-5,
      "the query path's own weight must be the bare normalized value: {q}"
    );
    assert!((p - q * tf).abs() < 1e-5, "{p} vs {q} * {tf}");
  }
}