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

use tempfile::tempdir;

use super::*;

fn write(dir: &Path, name: &str, body: &str) {
  let path = dir.join(name);
  fs::create_dir_all(path.parent().unwrap()).unwrap();
  fs::write(path, body).unwrap();
}

#[test]
fn discovers_single_safetensors_and_optional_files() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"bert"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "tokenizer_config.json", "{}");
  write(dir.path(), "model.safetensors", "");
  write(
    dir.path(),
    "1_Pooling/config.json",
    r#"{"pooling_mode_cls_token":true}"#,
  );
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(files.weights, vec![dir.path().join("model.safetensors")]);
  assert!(files.pooling_config.is_some());
  assert!(files.modules.is_none());
  assert!(files.tokenizer_config.is_some());
  assert_eq!(read_architecture(&files.config).unwrap(), "bert");
  assert_eq!(
    read_pooling(files.pooling_config.as_ref().unwrap()).unwrap(),
    Pooling::Cls
  );
}

#[test]
fn discovers_sharded_weights_in_order() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"bert"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(
    dir.path(),
    "model.safetensors.index.json",
    r#"{"weight_map":{"a":"model-00002-of-00002.safetensors","b":"model-00001-of-00002.safetensors","c":"model-00001-of-00002.safetensors"}}"#,
  );
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(
    files.weights,
    vec![
      dir.path().join("model-00001-of-00002.safetensors"),
      dir.path().join("model-00002-of-00002.safetensors")
    ]
  );
}

#[test]
fn discovers_a_torch_pickle_when_there_is_no_safetensors() {
  let dir = tempdir().unwrap();
  write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
  write(dir.path(), "tokenizer.json", "{}");
  write(dir.path(), "pytorch_model.bin", "");
  let files = discover(dir.path(), Pipeline::Dense).unwrap();
  assert_eq!(files.weights, vec![dir.path().join("pytorch_model.bin")]);
}

mod errors;
mod pipeline;
mod shape;