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

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use serde::Deserialize;
use serde::Serialize;

use crate::error::Result;

/// `get_slopes_power_of_2` in closed form: `slope_i = 2^(-8(i+1)/n)`.
fn power_of_two_slopes(n: usize) -> Vec<f64> {
  let ratio = 2f64.powf(-8.0 / n as f64);
  (0..n).map(|i| ratio.powi(i as i32 + 1)).collect()
}

/// `_get_alibi_head_slopes`: past the largest power of two `c`, the heads take every second slope of the `2 * c`
/// sequence, so the order is not monotonic, matching the checkpoint's own `alibi` buffer.
pub fn head_slopes(n: usize) -> Vec<f64> {
  if n.is_power_of_two() {
    return power_of_two_slopes(n);
  }
  let c = n.next_power_of_two() / 2;
  let mut slopes = power_of_two_slopes(c);
  slopes.extend(
    power_of_two_slopes(2 * c)
      .into_iter()
      .step_by(2)
      .take(n - c),
  );
  slopes
}

/// The ALiBi slope formula of a JinaBERT checkpoint whose head count is not a power of two; the variants differ only
/// in the heads past the largest power of two.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum AlibiSlopeVariant {
  /// The shared JinaBERT implementation's slopes, for every embedding checkpoint.
  Canonical,
  /// Those slopes with the tail heads halved, as `jinaai/jina-reranker-v1-tiny-en`'s bundled code computes them.
  HalvedInterpolatedTail,
}

/// ALiBi as `JinaBertEncoder.rebuild_alibi_tensor` defines it: `alibi[0, h, i, j] = -slope_h * |i - j|`.
/// Built per forward at the batch's length; the reference's cached buffer at 8192 positions is 3.2 GB in F32.
pub struct AlibiBias {
  slopes: Vec<f64>,
}

impl AlibiBias {
  /// `variant` matters only when `num_heads` is not a power of two; it halves the slopes past `head_slopes`' `c`.
  pub fn new(num_heads: usize, variant: AlibiSlopeVariant) -> Self {
    let mut slopes = head_slopes(num_heads);
    if variant == AlibiSlopeVariant::HalvedInterpolatedTail && !num_heads.is_power_of_two() {
      let tail_start = num_heads.next_power_of_two() / 2;
      for slope in &mut slopes[tail_start..] {
        *slope /= 2.0;
      }
    }
    Self { slopes }
  }

  /// `(1, heads, seq, seq)` in F32, as the reference builds it; `AttentionMask::padding_with_bias` casts it.
  pub fn forward(&self, seq: usize, device: &Device) -> Result<Tensor> {
    let positions = Tensor::arange(0u32, seq as u32, device)?.to_dtype(DType::F32)?;
    let distance = positions
      .reshape((seq, 1))?
      .broadcast_sub(&positions.reshape((1, seq))?)?
      .abs()?;
    let negated: Vec<f32> = self.slopes.iter().map(|s| -(*s as f32)).collect();
    let slopes = Tensor::from_vec(negated, (self.slopes.len(), 1, 1), device)?;
    Ok(
      distance
        .broadcast_mul(&slopes)?
        .unsqueeze(0)?
        .contiguous()?,
    )
  }
}

#[cfg(test)]
mod tests;