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

use tempfile::tempdir;

use super::*;

/// FlagEmbedding's own loader hides a missing file behind an info log and a randomly initialized head (`runner.py:90-97`).
#[test]
fn a_missing_colbert_linear_file_is_a_named_load_error() {
  let dir = tempdir().unwrap();
  let missing = dir.path().join("colbert_linear.pt");
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);

  let err =
    M3Colbert::load(xlm_config(tiny_config(4)), vb, &missing, conventions(false)).unwrap_err();
  assert!(
    matches!(err, CanddingError::MissingFile(ref p) if p == &missing),
    "{err}"
  );
  assert!(err.to_string().contains("colbert_linear.pt"), "{err}");
}

#[test]
fn from_raw_head_carries_keeps_cls_row_and_normalize_onto_the_struct() {
  let weight = Tensor::eye(4, DType::F32, &Device::Cpu).unwrap();
  let bias = Tensor::zeros(4, DType::F32, &Device::Cpu).unwrap();
  let config = xlm_config(tiny_config(4));
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let mut wanted = conventions(false);
  wanted.normalize = false;
  let model = M3Colbert::from_raw_head(config, vb, weight, Some(bias), wanted).unwrap();
  assert!(
    !model.keeps_cls_row,
    "keeps_cls_row must carry through from_raw_head"
  );
  assert!(
    !model.normalize,
    "normalize must carry through from_raw_head"
  );
}

/// `colbert_linear.pt` holds f16 tensors; without the cast, `forward` would fail on an f32 by f16 matmul.
#[test]
fn the_head_is_cast_from_f16_to_the_runs_dtype() {
  let weight = Tensor::eye(4, DType::F32, &Device::Cpu)
    .unwrap()
    .to_dtype(DType::F16)
    .unwrap();
  let bias = Tensor::zeros(4, DType::F32, &Device::Cpu)
    .unwrap()
    .to_dtype(DType::F16)
    .unwrap();
  assert_eq!(weight.dtype(), DType::F16);

  let config = xlm_config(tiny_config(4));
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = M3Colbert::from_raw_head(config, vb, weight, Some(bias), conventions(false)).unwrap();
  assert_eq!(model.weight.dtype(), DType::F32);
  assert_eq!(model.bias.as_ref().unwrap().dtype(), DType::F32);

  let input_ids = Tensor::from_vec(vec![4u32, 5, 6, 7], (1, 4), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32; 4], (1, 4), &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: &[4],
  };
  model.forward(&input).unwrap();
}

#[test]
fn dim_is_read_from_the_weights_own_shape() {
  let weight = Tensor::from_vec(
    (0..12).map(|v| v as f32 * 0.1).collect(),
    (3, 4),
    &Device::Cpu,
  )
  .unwrap();
  let bias = Tensor::zeros(3, DType::F32, &Device::Cpu).unwrap();
  let config = xlm_config(tiny_config(4));
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = M3Colbert::from_raw_head(config, vb, weight, Some(bias), conventions(false)).unwrap();
  assert_eq!(model.dim(), 3);

  let input_ids = Tensor::from_vec(vec![4u32, 5, 6, 7], (1, 4), &Device::Cpu).unwrap();
  let attention_mask = Tensor::from_vec(vec![1u32; 4], (1, 4), &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: &[4],
  };
  let rows = model.forward(&input).unwrap();
  assert_eq!(rows[0].dim(), 3);
}