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

use candle_core::Device;
use candle_core::Tensor;

use super::*;
use crate::models::Encoder;

/// Zeros at whatever hidden width it was built with: the guards read only shapes.
struct StubEncoder(usize);

impl Encoder for StubEncoder {
  fn forward(&self, input: &EncoderInput) -> Result<Tensor> {
    let (batch, seq) = input.input_ids.dims2()?;
    Ok(Tensor::zeros(
      (batch, seq, self.0),
      DType::F32,
      input.input_ids.device(),
    )?)
  }

  fn hidden_size(&self) -> usize {
    self.0
  }

  fn max_position_embeddings(&self) -> usize {
    16
  }

  fn weights_dtype(&self) -> DType {
    DType::F32
  }
}

/// A matrix large enough to reach `index::SHIFT` (32000 words × 4) is refused at exactly that boundary.
#[test]
fn a_matrix_large_enough_to_collide_with_the_hashed_range_is_refused() {
  let words: Vec<String> = (0..32_000).map(|i| format!("word{i}")).collect();
  let word_refs: Vec<&str> = words.iter().map(String::as_str).collect();
  let matrix = Tensor::zeros((32_001, 2, 4), DType::F32, &Device::Cpu).unwrap();
  // `MiniCoilSparse` is not `Debug`, so this matches rather than calling `unwrap_err`.
  match MiniCoilSparse::new(
    Box::new(StubEncoder(2)),
    matrix,
    vocab(&word_refs),
    stopwords(&[]),
    token_text(&["cat"]),
  ) {
    Err(err) => assert!(
      err.to_string().contains("128000"),
      "expected the compiled-in SHIFT boundary in the error, got: {err}"
    ),
    Ok(_) => panic!("a matrix this large must be refused"),
  }
}

/// One row narrower, its largest dense index is `31999 * 4 + 3 == 127999`, one short of `SHIFT`.
#[test]
fn a_matrix_exactly_at_the_collision_boundary_is_accepted() {
  let words: Vec<String> = (0..31_999).map(|i| format!("word{i}")).collect();
  let word_refs: Vec<&str> = words.iter().map(String::as_str).collect();
  let matrix = Tensor::zeros((32_000, 2, 4), DType::F32, &Device::Cpu).unwrap();
  assert!(
    MiniCoilSparse::new(
      Box::new(StubEncoder(2)),
      matrix,
      vocab(&word_refs),
      stopwords(&[]),
      token_text(&["cat"]),
    )
    .is_ok()
  );
}

/// A mismatched encoder width would otherwise make `Iterator::zip` drop the back half of every hidden state silently.
#[test]
fn a_hidden_width_disagreeing_with_the_matrix_is_refused() {
  let matrix = Tensor::from_vec(vec![0f32; 2 * 3 * 4], (2, 3, 4), &Device::Cpu).unwrap();
  let model = MiniCoilSparse::new(
    Box::new(StubEncoder(2)),
    matrix,
    vocab(&["cat"]),
    stopwords(&[]),
    token_text(&["cat"]),
  )
  .unwrap();

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

  let err = model.forward(&input).unwrap_err();
  let message = err.to_string();
  assert!(
    message.contains('2') && message.contains('3'),
    "expected both widths named in the error, got: {message}"
  );
}