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

//! The bytes a weight file's tensors take once loaded, read without loading any: a safetensors header's shapes, or
//! a pickle's or `.npy` file's size scaled from the dtype it stores.

use std::collections::HashMap;
use std::fs;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use candle_core::DType;
use serde::Deserialize;

use super::npy;
use crate::error::CanddingError;
use crate::error::Result;

/// safetensors' own cap, so a corrupt length cannot ask for a huge allocation.
const MAX_HEADER_BYTES: u64 = 100_000_000;

#[derive(Deserialize)]
struct TensorEntry {
  shape: Vec<u64>,
}

/// The bytes `path`'s tensors take at `dtype`; a file that is neither safetensors nor `.npy` is a torch pickle.
pub(crate) fn resident_bytes(path: &Path, dtype: DType) -> Result<u64> {
  match path.extension().and_then(|ext| ext.to_str()) {
    Some("safetensors") => {
      Ok(safetensors_elements(path)?.saturating_mul(dtype.size_in_bytes() as u64))
    }
    Some("npy") => Ok(scale(
      fs::metadata(path)?.len(),
      npy::stored_dtype(path)?,
      dtype,
    )),
    _ => Ok(scale(fs::metadata(path)?.len(), pickle_dtype(path)?, dtype)),
  }
}

/// `bytes` of elements stored in `from`, as the bytes the same elements take in `to`.
pub(crate) fn scale(bytes: u64, from: DType, to: DType) -> u64 {
  let scaled = u128::from(bytes) * to.size_in_bytes() as u128 / from.size_in_bytes() as u128;
  u64::try_from(scaled).unwrap_or(u64::MAX)
}

fn safetensors_elements(path: &Path) -> Result<u64> {
  let mut file = File::open(path)?;
  let mut length = [0u8; 8];
  file
    .read_exact(&mut length)
    .map_err(|e| unreadable(path, &e.to_string()))?;
  let length = u64::from_le_bytes(length);
  if length > MAX_HEADER_BYTES {
    return Err(unreadable(
      path,
      &format!("its header claims {length} bytes"),
    ));
  }
  let mut header = vec![0u8; length as usize];
  file
    .read_exact(&mut header)
    .map_err(|e| unreadable(path, &e.to_string()))?;
  let mut entries: HashMap<String, serde_json::Value> =
    serde_json::from_slice(&header).map_err(|e| unreadable(path, &e.to_string()))?;
  entries.remove("__metadata__");
  entries.into_values().try_fold(0u64, |total, entry| {
    let entry: TensorEntry =
      serde_json::from_value(entry).map_err(|e| unreadable(path, &e.to_string()))?;
    let elements = entry
      .shape
      .iter()
      .fold(1u64, |n, &dim| n.saturating_mul(dim));
    Ok(total.saturating_add(elements))
  })
}

/// The dtype of a pickle's largest tensor, from the tensor list candle reads ahead of any tensor data.
fn pickle_dtype(path: &Path) -> Result<DType> {
  let infos = candle_core::pickle::read_pth_tensor_info(path, false, None)?;
  let largest = infos
    .iter()
    .max_by_key(|info| info.layout.shape().elem_count());
  Ok(largest.map_or(DType::F32, |info| info.dtype))
}

fn unreadable(path: &Path, reason: &str) -> CanddingError {
  candle_core::Error::Msg(format!(
    "`{}` has no readable safetensors header: {reason}",
    path.display()
  ))
  .into()
}

#[cfg(test)]
mod tests;