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

use std::path::Path;
use std::path::PathBuf;

use candle_core::DType;
use tempfile::tempdir;

use super::*;

/// Not padded to numpy's 64-byte alignment, so the reader must rely on the header-length field alone.
fn npy_bytes(descr: &str, fortran_order: bool, shape: &[usize], data: &[u8]) -> Vec<u8> {
  let shape_str = match shape {
    [] => "()".to_string(),
    [only] => format!("({only},)"),
    many => format!(
      "({})",
      many
        .iter()
        .map(usize::to_string)
        .collect::<Vec<_>>()
        .join(", ")
    ),
  };
  let dict = format!(
    "{{'descr': '{descr}', 'fortran_order': {}, 'shape': {shape_str}, }}\n",
    if fortran_order { "True" } else { "False" }
  );
  let mut bytes = Vec::new();
  bytes.extend_from_slice(&MAGIC);
  bytes.extend_from_slice(&[1, 0]);
  bytes.extend_from_slice(&(dict.len() as u16).to_le_bytes());
  bytes.extend_from_slice(dict.as_bytes());
  bytes.extend_from_slice(data);
  bytes
}

fn write_file(dir: &Path, name: &str, bytes: &[u8]) -> PathBuf {
  let path = dir.join(name);
  std::fs::write(&path, bytes).unwrap();
  path
}

fn le_f32(values: &[f32]) -> Vec<u8> {
  values.iter().flat_map(|v| v.to_le_bytes()).collect()
}

#[test]
fn round_trip_through_a_file_this_test_writes_itself() {
  let dir = tempdir().unwrap();
  let values = [1.0f32, -2.5, 3.25, 4.0, -5.5, 6.75];
  let bytes = npy_bytes("<f4", false, &[2, 3], &le_f32(&values));
  let path = write_file(dir.path(), "round_trip.npy", &bytes);

  let tensor = read(&path).unwrap();
  assert_eq!(tensor.dims(), &[2, 3]);
  assert_eq!(tensor.dtype(), DType::F32);
  assert_eq!(
    tensor.to_vec2::<f32>().unwrap(),
    vec![vec![1.0, -2.5, 3.25], vec![4.0, -5.5, 6.75]]
  );
}

#[test]
fn round_trip_preserves_a_single_byte_integer_dtype() {
  let dir = tempdir().unwrap();
  let bytes = npy_bytes("|u1", false, &[4], &[0u8, 1, 254, 255]);
  let path = write_file(dir.path(), "bytes.npy", &bytes);

  let tensor = read(&path).unwrap();
  assert_eq!(tensor.dtype(), DType::U8);
  assert_eq!(tensor.to_vec1::<u8>().unwrap(), vec![0, 1, 254, 255]);
}

/// 156_672_128 bytes minus a 128-byte header, over 4 bytes and 512 * 4, is exactly the declared 19_125 words.
#[test]
fn minicoil_file_byte_count_matches_its_declared_shape() {
  const TOTAL_BYTES: usize = 156_672_128;
  const STANDARD_HEADER_BYTES: usize = 128;
  const INPUT_DIM: usize = 512;
  const OUTPUT_DIM: usize = 4;
  const F32_BYTES: usize = 4;

  let data_bytes = TOTAL_BYTES - STANDARD_HEADER_BYTES;
  assert_eq!(data_bytes, 156_672_000);
  let floats = data_bytes / F32_BYTES;
  assert_eq!(floats, 39_168_000);
  assert_eq!(floats % (INPUT_DIM * OUTPUT_DIM), 0);
  assert_eq!(floats / (INPUT_DIM * OUTPUT_DIM), 19_125);
}

mod errors;