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

use std::collections::HashMap;

use super::*;

/// `projection_in: 4` matches `tiny_config(4)`; `projection_out` mirrors the tensor map's `out_features`.
fn conventions(keeps_cls_row: bool, projection_out: usize) -> MultiVectorConventions {
  MultiVectorConventions {
    keeps_cls_row,
    projection_in: 4,
    projection_out,
    projection_bias: false,
    normalize: true,
    attends_to_mask_tokens: false,
    position_scheme: None,
  }
}

fn tensor_map(hidden_size: usize, out_features: usize) -> HashMap<String, Tensor> {
  let mut tensors = HashMap::new();
  tensors.insert(
    "linear.weight".to_string(),
    Tensor::zeros((out_features, hidden_size), DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "embeddings.word_embeddings.weight".to_string(),
    Tensor::zeros((64, hidden_size), DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "embeddings.position_embeddings.weight".to_string(),
    Tensor::zeros((32, hidden_size), DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "embeddings.token_type_embeddings.weight".to_string(),
    Tensor::zeros((2, hidden_size), DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "embeddings.LayerNorm.weight".to_string(),
    Tensor::ones(hidden_size, DType::F32, &Device::Cpu).unwrap(),
  );
  tensors.insert(
    "embeddings.LayerNorm.bias".to_string(),
    Tensor::zeros(hidden_size, DType::F32, &Device::Cpu).unwrap(),
  );
  tensors
}

/// A `VarMap` auto-creates a missing tensor, so only `VarBuilder::from_tensors` proves this shape check fires.
#[test]
fn load_reads_the_real_root_linear_weight_and_its_shape() {
  let tensors = tensor_map(4, 3);
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let model =
    StanfordColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 3)).unwrap();
  assert_eq!(model.dim(), 3);
}

#[test]
fn a_projection_out_that_disagrees_with_the_real_tensor_shape_is_refused() {
  let tensors = tensor_map(4, 3);
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  // `StanfordColbert` is not `Debug`, so this matches the `Result` rather than calling `unwrap_err`.
  match StanfordColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 5)) {
    Ok(_) => panic!("a projection_out that disagrees with the real tensor shape must be refused"),
    Err(err) => assert!(
      matches!(err, crate::error::CanddingError::Candle(_)),
      "{err}"
    ),
  }
}

#[test]
fn a_missing_root_linear_weight_is_refused() {
  let mut tensors = tensor_map(4, 3);
  tensors.remove("linear.weight");
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  match StanfordColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 3)) {
    Ok(_) => panic!("a missing root linear.weight must be refused"),
    Err(err) => assert!(
      matches!(err, crate::error::CanddingError::Candle(_)),
      "{err}"
    ),
  }
}

#[test]
fn load_carries_keeps_cls_row_and_normalize_onto_the_struct() {
  let tensors = tensor_map(4, 3);
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let mut wanted = conventions(false, 3);
  wanted.normalize = false;
  let model = StanfordColbert::load(tiny_config(4), vb, HashSet::new(), wanted).unwrap();
  assert!(
    !model.keeps_cls_row,
    "keeps_cls_row must carry through load"
  );
  assert!(!model.normalize, "normalize must carry through load");
}

#[test]
fn projection_bias_gates_whether_linear_bias_is_read() {
  let tensors = tensor_map(4, 3);
  let vb = VarBuilder::from_tensors(tensors.clone(), DType::F32, &Device::Cpu);
  assert!(
    StanfordColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 3)).is_ok(),
    "projection_bias:false must not try to read a linear.bias that is not there"
  );
  let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
  let mut wants_bias = conventions(true, 3);
  wants_bias.projection_bias = true;
  match StanfordColbert::load(tiny_config(4), vb, HashSet::new(), wants_bias) {
    Ok(_) => panic!("projection_bias:true must fail when linear.bias is not published"),
    Err(err) => assert!(
      matches!(err, crate::error::CanddingError::Candle(_)),
      "{err}"
    ),
  }
}