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

//! `linear_head.bin` is a real `torch.save` archive shaped like `colbert_linear.pt`, from `scripts/pickle_fixture.py`.

use super::*;

fn fixture(name: &str) -> std::path::PathBuf {
  std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
    .join("tests")
    .join("fixtures")
    .join("pickle")
    .join(name)
}

#[test]
fn load_reads_weight_and_bias_from_a_real_archive() {
  let config = xlm_config(tiny_config(4));
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let model = M3Colbert::load(config, vb, &fixture("linear_head.bin"), conventions(false)).unwrap();

  assert_eq!(
    model
      .weight
      .to_dtype(DType::F32)
      .unwrap()
      .to_vec2::<f32>()
      .unwrap(),
    vec![vec![1.5f32, -2.0, 0.5, 0.25]]
  );
  let bias = model
    .bias
    .as_ref()
    .expect("projection_bias: true must read bias");
  assert_eq!(
    bias.to_dtype(DType::F32).unwrap().to_vec1::<f32>().unwrap(),
    vec![0.75f32]
  );

  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 load_does_not_read_bias_when_projection_bias_is_false() {
  let config = xlm_config(tiny_config(4));
  let varmap = VarMap::new();
  let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
  let mut conv = conventions(false);
  conv.projection_bias = false;

  let model = M3Colbert::load(config, vb, &fixture("linear_head.bin"), conv).unwrap();

  assert!(
    model.bias.is_none(),
    "projection_bias: false must not read bias even though linear_head.bin has one"
  );
}