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

//! Steps applied to pooled vectors: truncation, normalization, extraction.

use candle_core::DType;
use candle_core::Tensor;

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

/// L2-normalizes each row; the norm is clamped at 1e-12 to avoid dividing by zero.
pub fn normalize(xs: &Tensor) -> Result<Tensor> {
  let norms = xs.sqr()?.sum_keepdim(1)?.sqrt()?.maximum(1e-12)?;
  Ok(xs.broadcast_div(&norms)?)
}

/// `layer_norm` over the last dimension with no weight or bias, as `nomic-embed-text-v1.5`'s card applies it
/// before truncation; it changes the full-width vector's direction too.
pub fn layer_norm(xs: &Tensor, eps: f64) -> Result<Tensor> {
  let mean = xs.mean_keepdim(1)?;
  let centered = xs.broadcast_sub(&mean)?;
  let var = centered.sqr()?.mean_keepdim(1)?;
  Ok(centered.broadcast_div(&(var + eps)?.sqrt()?)?)
}

/// Keeps the first `dim` columns, the Matryoshka prefix.
pub fn truncate_dim(xs: &Tensor, dim: usize) -> Result<Tensor> {
  let full = xs.dim(1)?;
  if dim == 0 || dim > full {
    return Err(CanddingError::InvalidDim {
      requested: dim,
      dim: full,
    });
  }
  Ok(xs.narrow(1, 0, dim)?)
}

pub fn to_rows(xs: &Tensor) -> Result<Vec<Vec<f32>>> {
  Ok(xs.to_dtype(DType::F32)?.to_vec2::<f32>()?)
}

#[cfg(test)]
mod tests {
  use candle_core::Device;

  use super::*;

  #[test]
  fn normalize_gives_unit_rows() {
    let xs = Tensor::new(&[[3f32, 4.0], [0.0, 2.0]], &Device::Cpu).unwrap();
    let rows = to_rows(&normalize(&xs).unwrap()).unwrap();
    assert_eq!(rows, vec![vec![0.6, 0.8], vec![0.0, 1.0]]);
  }

  #[test]
  fn truncate_keeps_prefix_and_rejects_bad_dims() {
    let xs = Tensor::new(&[[1f32, 2.0, 3.0]], &Device::Cpu).unwrap();
    assert_eq!(
      to_rows(&truncate_dim(&xs, 2).unwrap()).unwrap(),
      vec![vec![1.0, 2.0]]
    );
    assert!(matches!(
      truncate_dim(&xs, 4).unwrap_err(),
      CanddingError::InvalidDim {
        requested: 4,
        dim: 3
      }
    ));
    assert!(truncate_dim(&xs, 0).is_err());
  }

  /// Row `[1, 2, 3, 6]` (mean 3, biased variance 3.5) gives `[-1.0690, -0.5345, 0, 1.6036]`; uncentered, it would sum to about 3.39.
  #[test]
  fn layer_norm_removes_the_row_mean_and_scales_by_the_biased_std() {
    let xs = Tensor::new(&[[1f32, 2.0, 3.0, 6.0]], &Device::Cpu).unwrap();
    let out = to_rows(&layer_norm(&xs, 1e-5).unwrap()).unwrap();
    let expected = [-1.069_045, -0.534_522, 0.0, 1.603_567];
    for (o, e) in out[0].iter().zip(expected) {
      assert!((o - e).abs() < 1e-4, "{out:?}");
    }
    assert!(
      out[0].iter().sum::<f32>().abs() < 1e-5,
      "the row mean is removed"
    );
  }
}