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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

//! The sparse embedding output: a row of `(index, weight)` pairs over a model's own output space,
//! disjoint from `Encoder`'s dense `(batch, seq, hidden)` tensor.

use std::borrow::Cow;
use std::cmp::Ordering;
use std::collections::BTreeMap;

/// A sparse row of `(index, weight)` pairs over a head-specific index space; a weight may be negative.
/// Every head and `from_pairs` produce strictly ascending indices, but `dot` accepts any order.
#[derive(Clone, Debug, PartialEq)]
pub struct SparseEmbedding {
  pub indices: Vec<u32>,
  pub values: Vec<f32>,
}

impl SparseEmbedding {
  /// Builds a row from `(index, weight)` pairs, keeping the largest weight of a repeated index and dropping exact
  /// zeros; negative weights stay, and the result is ascending by index.
  pub fn from_pairs(pairs: impl IntoIterator<Item = (u32, f32)>) -> Self {
    let mut folded: BTreeMap<u32, f32> = BTreeMap::new();
    for (index, value) in pairs {
      if value == 0.0 {
        continue;
      }
      let kept = folded.entry(index).or_insert(value);
      if value > *kept {
        *kept = value;
      }
    }
    let (indices, values): (Vec<u32>, Vec<f32>) = folded.into_iter().unzip();
    Self { indices, values }
  }

  /// The number of entries.
  pub fn len(&self) -> usize {
    self.indices.len()
  }

  /// Whether the row has no entries.
  pub fn is_empty(&self) -> bool {
    self.indices.is_empty()
  }

  /// The dot product over the shared indices of both rows as `from_pairs` would build them, so any index order works.
  /// An index without a value at the same position, or a value without an index, is ignored.
  pub fn dot(&self, other: &Self) -> f32 {
    let (a, b) = (self.ascending(), other.ascending());
    let mut sum = 0.0;
    let (mut i, mut j) = (0, 0);
    while i < a.indices.len() && j < b.indices.len() {
      match a.indices[i].cmp(&b.indices[j]) {
        Ordering::Less => i += 1,
        Ordering::Greater => j += 1,
        Ordering::Equal => {
          sum += a.values[i] * b.values[j];
          i += 1;
          j += 1;
        }
      }
    }
    sum
  }

  /// This row when it is already strictly ascending with one value per index, else `from_pairs` of its pairs.
  fn ascending(&self) -> Cow<'_, Self> {
    if self.indices.len() == self.values.len() && self.indices.is_sorted_by(|a, b| a < b) {
      Cow::Borrowed(self)
    } else {
      let pairs = self
        .indices
        .iter()
        .copied()
        .zip(self.values.iter().copied());
      Cow::Owned(Self::from_pairs(pairs))
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  /// Duplicates fold by maximum, only exact zeros drop, and indices ascend so a dot product is a merge;
  /// a negative value stays, as `minicoil`'s signed `tanh` output needs.
  #[test]
  fn from_pairs_folds_duplicates_by_maximum_and_drops_only_exact_zero() {
    let s =
      SparseEmbedding::from_pairs([(7u32, 0.5f32), (3, 0.25), (7, 1.5), (9, 0.0), (11, -0.5)]);
    assert_eq!(s.indices, vec![3, 7, 11]);
    assert_eq!(s.values, vec![0.25, 1.5, -0.5]);
    assert!(SparseEmbedding::from_pairs([]).is_empty());
  }

  #[test]
  fn dot_multiplies_the_shared_keys_only() {
    let a = SparseEmbedding::from_pairs([(1u32, 2.0f32), (5, 3.0)]);
    let b = SparseEmbedding::from_pairs([(5u32, 4.0f32), (9, 100.0)]);
    assert_eq!(a.dot(&b), 12.0);
    assert_eq!(a.dot(&SparseEmbedding::from_pairs([])), 0.0);
  }

  #[test]
  fn dot_accepts_indices_in_any_order() {
    let unsorted = SparseEmbedding {
      indices: vec![5, 1],
      values: vec![1.0, 2.0],
    };
    let sorted = SparseEmbedding::from_pairs([(1u32, 1.0f32), (5, 1.0)]);
    assert_eq!(unsorted.dot(&sorted), 3.0);
    assert_eq!(sorted.dot(&unsorted), 3.0);
  }

  #[test]
  fn dot_ignores_an_index_without_a_value() {
    let ragged = SparseEmbedding {
      indices: vec![1, 5],
      values: vec![1.0],
    };
    let sorted = SparseEmbedding::from_pairs([(1u32, 1.0f32), (5, 1.0)]);
    assert_eq!(ragged.dot(&sorted), 1.0);
    assert_eq!(sorted.dot(&ragged), 1.0);
  }

  #[test]
  fn dot_folds_a_repeated_index_by_maximum() {
    let repeated = SparseEmbedding {
      indices: vec![1, 1],
      values: vec![1.0, 2.0],
    };
    let one = SparseEmbedding::from_pairs([(1u32, 1.0f32)]);
    assert_eq!(repeated.dot(&one), 2.0);
  }
}