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

//! Picks the rotary base and length for NomicBERT: base 1000, dynamic NTK past the trained length when configured.

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

use super::config::NomicBertConfig;
use crate::error::Result;
use crate::layers::rope::RotaryEmbedding;

/// Dynamic-NTK base, in f64: `base * ((factor * seq / max_trained) - (factor - 1)) ^ (dim / (dim - 2))`.
/// `base` unchanged at or below `max_trained`.
pub fn dynamic_ntk_base(base: f64, factor: f64, max_trained: usize, seq: usize, dim: usize) -> f64 {
  if seq <= max_trained {
    return base;
  }
  let scaled = (factor * seq as f64 / max_trained as f64) - (factor - 1.0);
  base * scaled.powf(dim as f64 / (dim as f64 - 2.0))
}

/// `cached` covers `0..max_trained_positions`; a longer batch gets a fresh table instead of growing it.
pub struct NomicRotary {
  base: f64,
  factor: Option<f64>,
  max_trained: usize,
  head_dim: usize,
  dtype: DType,
  device: Device,
  cached: RotaryEmbedding,
}

impl NomicRotary {
  pub fn new(config: &NomicBertConfig, dtype: DType, device: &Device) -> Result<Self> {
    let head_dim = config.head_dim();
    let max_trained = config.max_trained_positions;
    let base = config.rotary_emb_base;
    let cached = RotaryEmbedding::new(head_dim, max_trained, base, dtype, device)?;
    Ok(Self {
      base,
      factor: config.scaling_factor(),
      max_trained,
      head_dim,
      dtype,
      device: device.clone(),
      cached,
    })
  }

  /// The cached table up to the trained length; past it, a fresh one at the dynamic-NTK (or plain, for `v1.5`) base.
  /// Unlike the reference it rescales by the batch's padded length and caches nothing; no fixture goes that far.
  pub fn for_length(&self, seq: usize) -> Result<Table<'_>> {
    if seq <= self.max_trained {
      return Ok(Table::Cached(&self.cached));
    }
    let base = self.factor.map_or(self.base, |factor| {
      dynamic_ntk_base(self.base, factor, self.max_trained, seq, self.head_dim)
    });
    Ok(Table::Fresh(RotaryEmbedding::new(
      self.head_dim,
      seq,
      base,
      self.dtype,
      &self.device,
    )?))
  }
}

pub enum Table<'a> {
  Cached(&'a RotaryEmbedding),
  Fresh(RotaryEmbedding),
}

impl Table<'_> {
  pub(super) fn as_rope(&self) -> &RotaryEmbedding {
    match self {
      Table::Cached(rope) => rope,
      Table::Fresh(rope) => rope,
    }
  }
}

#[cfg(test)]
mod tests;