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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

//! Parses the npy header dictionary into a dtype, its byte width and the shape; Fortran order is refused.

use std::path::Path;

use candle_core::DType;

use crate::error::CanddingError;
use crate::error::Result;

/// A `descr` suffix after the byte-order character, its candle dtype and width; single-letter kind codes are refused.
struct DtypeEntry {
  code: &'static str,
  dtype: DType,
  itemsize: usize,
}

const SUPPORTED_DTYPES: &[DtypeEntry] = &[
  DtypeEntry {
    code: "f4",
    dtype: DType::F32,
    itemsize: 4,
  },
  DtypeEntry {
    code: "f8",
    dtype: DType::F64,
    itemsize: 8,
  },
  DtypeEntry {
    code: "i2",
    dtype: DType::I16,
    itemsize: 2,
  },
  DtypeEntry {
    code: "i4",
    dtype: DType::I32,
    itemsize: 4,
  },
  DtypeEntry {
    code: "i8",
    dtype: DType::I64,
    itemsize: 8,
  },
  DtypeEntry {
    code: "u1",
    dtype: DType::U8,
    itemsize: 1,
  },
  DtypeEntry {
    code: "u4",
    dtype: DType::U32,
    itemsize: 4,
  },
];

pub(super) struct Header {
  pub(super) dtype: DType,
  pub(super) itemsize: usize,
  pub(super) shape: Vec<usize>,
}

/// Parses the dictionary body, e.g. `{'descr': '<f4', 'fortran_order': False, 'shape': (19125, 512, 4), }`;
/// it is always that flat shape, so splitting on top-level commas and the first `:` suffices.
pub(super) fn parse_header(path: &Path, dict: &str) -> Result<Header> {
  let mut descr = None;
  let mut fortran_order = None;
  let mut shape = None;
  for field in split_top_level(dict) {
    let Some((key, value)) = field.split_once(':') else {
      continue;
    };
    let key = key.trim().trim_matches('\'');
    let value = value.trim();
    match key {
      "descr" => descr = Some(value.trim_matches('\'').to_string()),
      "fortran_order" => fortran_order = Some(value == "True"),
      "shape" => {
        shape = Some(parse_shape(value).map_err(|reason| super::header_error(path, &reason))?)
      }
      _ => {}
    }
  }
  let descr =
    descr.ok_or_else(|| super::header_error(path, "header dictionary has no `descr` key"))?;
  let fortran_order = fortran_order
    .ok_or_else(|| super::header_error(path, "header dictionary has no `fortran_order` key"))?;
  let shape =
    shape.ok_or_else(|| super::header_error(path, "header dictionary has no `shape` key"))?;

  if fortran_order {
    return Err(CanddingError::NpyFortranOrder(path.to_path_buf()));
  }
  let (dtype, itemsize) = resolve_dtype(path, &descr)?;
  Ok(Header {
    dtype,
    itemsize,
    shape,
  })
}

/// `<` and `=` read as little-endian (every target is), `|` only for one-byte dtypes; `>` and a wider `|` are
/// refused rather than misread.
fn resolve_dtype(path: &Path, descr: &str) -> Result<(DType, usize)> {
  let mut chars = descr.chars();
  let order = chars.next().unwrap_or('\0');
  let code = chars.as_str();
  let Some(entry) = SUPPORTED_DTYPES.iter().find(|entry| entry.code == code) else {
    return Err(CanddingError::NpyUnsupportedDtype {
      file: path.to_path_buf(),
      descr: descr.to_string(),
    });
  };
  let order_ok = matches!(order, '<' | '=') || (order == '|' && entry.itemsize == 1);
  if !order_ok {
    return Err(CanddingError::NpyUnsupportedByteOrder {
      file: path.to_path_buf(),
      order,
      descr: descr.to_string(),
    });
  }
  Ok((entry.dtype, entry.itemsize))
}

/// Splits on top-level commas, not those inside the `shape` tuple, the only nesting a header has.
fn split_top_level(dict: &str) -> Vec<String> {
  let body = dict
    .trim()
    .trim_matches(|c: char| c == '{' || c == '}')
    .trim();
  let mut fields = Vec::new();
  let mut depth = 0i32;
  let mut start = 0usize;
  for (index, c) in body.char_indices() {
    match c {
      '(' => depth += 1,
      ')' => depth -= 1,
      ',' if depth == 0 => {
        fields.push(body[start..index].to_string());
        start = index + 1;
      }
      _ => {}
    }
  }
  if start < body.len() {
    fields.push(body[start..].to_string());
  }
  fields
}

/// `(19125, 512, 4)` or `()` for a scalar; the trailing comma numpy writes leaves an empty entry, dropped here.
fn parse_shape(value: &str) -> std::result::Result<Vec<usize>, String> {
  let inner = value.trim().trim_matches(|c| c == '(' || c == ')');
  inner
    .split(',')
    .map(str::trim)
    .filter(|entry| !entry.is_empty())
    .map(|entry| {
      entry
        .parse::<usize>()
        .map_err(|_| format!("shape entry `{entry}` is not a non-negative integer"))
    })
    .collect()
}