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

//! Refusing a load whose weights need more bytes than the device has available, before any of them is read.

use std::path::PathBuf;

use candle_core::DType;
use candle_core::Device;

use crate::device;
use crate::error::CanddingError;
use crate::error::Result;
use crate::hub;
use crate::hub::ModelSource;
use crate::registry::Backend;
use crate::weights::size;

/// A test hook: the budget in bytes, in place of the device's own figure.
const BUDGET_VAR: &str = "CANDDING_TEST_MEMORY_BUDGET";

/// The rule: weights fit when they need no more bytes than are available.
pub(crate) fn fits(needed_bytes: u64, available_bytes: u64) -> bool {
  needed_bytes <= available_bytes
}

/// Refuses a model with no snapshot on disk from its registry size, before the download. The size is read as
/// stored in `stored`, which under-counts a half-precision checkpoint registered as F32 until `check_files`.
pub(crate) fn check_download(
  model: &str,
  device: &Device,
  source: &ModelSource,
  size_gb: f32,
  stored: DType,
  dtype: DType,
) -> Result<()> {
  if size_gb <= 0.0 || hub::is_present(source) {
    return Ok(());
  }
  let download = (f64::from(size_gb) * 1e9) as u64;
  check(model, device, size::scale(download, stored, dtype))
}

/// Refuses weight files, each paired with the dtype its tensors load in, that cannot fit.
pub(crate) fn check_files(model: &str, device: &Device, files: &[(PathBuf, DType)]) -> Result<()> {
  let mut needed_bytes = 0u64;
  for (path, dtype) in files {
    needed_bytes = needed_bytes.saturating_add(size::resident_bytes(path, *dtype)?);
  }
  check(model, device, needed_bytes)
}

pub(crate) fn at(paths: &[PathBuf], dtype: DType) -> Vec<(PathBuf, DType)> {
  paths.iter().map(|path| (path.clone(), dtype)).collect()
}

fn check(model: &str, device: &Device, needed_bytes: u64) -> Result<()> {
  let injected = std::env::var(BUDGET_VAR)
    .ok()
    .and_then(|bytes| bytes.parse().ok());
  let available = injected.or_else(|| device::budget_bytes(device));
  verdict(model, Backend::of(device), needed_bytes, available)
}

/// No figure for the device means no limit.
fn verdict(model: &str, backend: Backend, needed_bytes: u64, available: Option<u64>) -> Result<()> {
  let Some(available_bytes) = available else {
    return Ok(());
  };
  tracing::debug!(
    model,
    backend = backend.name(),
    needed_bytes,
    available_bytes,
    "memory fit check"
  );
  if fits(needed_bytes, available_bytes) {
    return Ok(());
  }
  Err(CanddingError::InsufficientMemory {
    model: model.to_string(),
    backend,
    needed_bytes,
    available_bytes,
  })
}

#[cfg(test)]
mod tests;