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

use std::path::PathBuf;

use candle_core::Device;

use super::tiny;
use crate::tokenizer::PadToken;
use crate::tokenizer::TextTokenizer;

fn markers_path() -> PathBuf {
  PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer_markers.json")
}

/// `[MASK]` (id 18 here) is what `colbert-ir/colbertv2.0` and `answerdotai/answerai-colbert-small-v1` fill a query's tail with.
#[test]
fn fixed_length_padding_pads_every_row_to_the_requested_width_with_mask() {
  let tokenizer = TextTokenizer::from_file_fixed_length(
    &markers_path(),
    8,
    PadToken::Named("[MASK]".to_string()),
  )
  .unwrap();
  let batch = tokenizer.encode(&["hello"], &Device::Cpu).unwrap();
  assert_eq!(batch.input_ids.dims(), &[1, 8]);
  assert_eq!(
    batch.input_ids.to_vec2::<u32>().unwrap()[0],
    vec![2, 4, 3, 18, 18, 18, 18, 18],
    "3 real tokens (CLS, hello, SEP) then [MASK] filling the rest"
  );
  assert_eq!(
    batch.attention_mask.to_vec2::<u32>().unwrap()[0],
    vec![1, 1, 1, 0, 0, 0, 0, 0],
    "the [MASK]-filled tail is unattendable, exactly like ordinary padding"
  );
}

#[test]
fn an_unregistered_marker_shatters_instead_of_tokenizing_as_one_piece() {
  let tokenizer = tiny(16);
  let ids = tokenizer
    .encode(&["[unused0] hello"], &Device::Cpu)
    .unwrap()
    .input_ids
    .to_vec2::<u32>()
    .unwrap();
  assert!(
    !ids[0].contains(&16),
    "id 16 ([unused0]) must not appear from unregistered text: {:?}",
    ids[0]
  );
  assert!(
    ids[0].contains(&1),
    "the shattered pieces fall back to [UNK]: {:?}",
    ids[0]
  );
}

#[test]
fn a_registered_marker_tokenizes_as_one_piece_with_the_trailing_space_intact() {
  let mut tokenizer = TextTokenizer::from_file(&markers_path(), 16, PadToken::Id(0)).unwrap();
  tokenizer.register_special_tokens(&["[unused0]"]).unwrap();
  let ids = tokenizer
    .encode(&["[unused0] hello"], &Device::Cpu)
    .unwrap()
    .input_ids
    .to_vec2::<u32>()
    .unwrap();
  assert_eq!(
    ids[0],
    vec![2, 16, 4, 3],
    "CLS, [unused0], hello, SEP — one piece for the marker, not several"
  );
}

/// A registered marker tokenizes the same with or without the reference prompt's trailing space.
#[test]
fn a_registered_marker_tokenizes_identically_whether_or_not_the_following_text_had_its_own_space() {
  let mut tokenizer = TextTokenizer::from_file(&markers_path(), 16, PadToken::Id(0)).unwrap();
  tokenizer.register_special_tokens(&["[unused0]"]).unwrap();
  let spaced = tokenizer
    .encode(&["[unused0] hello"], &Device::Cpu)
    .unwrap()
    .input_ids
    .to_vec2::<u32>()
    .unwrap()[0]
    .clone();
  let unspaced = tokenizer
    .encode(&["[unused0]hello"], &Device::Cpu)
    .unwrap()
    .input_ids
    .to_vec2::<u32>()
    .unwrap()[0]
    .clone();
  assert_eq!(spaced, vec![2, 16, 4, 3]);
  assert_eq!(
    spaced, unspaced,
    "this crate's added-token matching does not need the space to recover [unused0] as one piece; \
     the pipeline still builds the spaced string to match the reference's own literal prompt"
  );
}