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

//! `pool_row`'s recognized-word half: average hidden states, gather the word's matrix slice, project, normalize.

use std::collections::HashMap;

use candle_core::Tensor;

use super::index;
use super::term_frequency;
use crate::error::Result;

/// The reference's norm floor, so a zero projection normalizes to zero, not `NaN`.
const NORM_FLOOR: f32 = 1e-8;

/// One resolved id's hidden-state sum, pooled flatly over its positions like `avg_by_vocab_ids`, not a mean of means.
pub(super) struct RecognizedGroup {
  token_sum: Vec<f32>,
  token_count: u32,
  /// Word occurrences for term frequency; `token_count` counts sub-token positions instead.
  pub(super) word_count: u32,
}

impl RecognizedGroup {
  pub(super) fn new(hidden_size: usize) -> Self {
    Self {
      token_sum: vec![0.0; hidden_size],
      token_count: 0,
      word_count: 0,
    }
  }

  pub(super) fn add_token(&mut self, hidden: &[f32]) {
    for (sum, value) in self.token_sum.iter_mut().zip(hidden) {
      *sum += value;
    }
    self.token_count += 1;
  }

  /// `token_count` is never 0: a group exists only after `add_token` ran.
  fn average(&self) -> Vec<f32> {
    let count = self.token_count as f32;
    self.token_sum.iter().map(|sum| sum / count).collect()
  }
}

/// `tanh` output over its L2 norm along axis 1, floored at `NORM_FLOOR`.
fn l2_normalize(x: &Tensor) -> Result<Tensor> {
  let floor = Tensor::new(&[NORM_FLOOR], x.device())?;
  let norm = x.sqr()?.sum_keepdim(1)?.sqrt()?.broadcast_maximum(&floor)?;
  Ok(x.broadcast_div(&norm)?)
}

/// One batched `index_select` and matmul over the distinct resolved words' `(hidden, embedding_size)` slices.
pub(super) fn pairs(
  groups: &HashMap<u32, RecognizedGroup>,
  matrix: &Tensor,
  doc_len: u32,
  uniform_weight: bool,
) -> Result<Vec<(u32, f32)>> {
  if groups.is_empty() {
    return Ok(Vec::new());
  }
  let device = matrix.device();
  let hidden_size = matrix.dims()[1];
  let word_ids: Vec<u32> = groups.keys().copied().collect();
  let mut flat_avg = Vec::with_capacity(word_ids.len() * hidden_size);
  for &word_id in &word_ids {
    flat_avg.extend(groups[&word_id].average());
  }
  let ids = Tensor::from_vec(word_ids.clone(), word_ids.len(), device)?;
  let avg = Tensor::from_vec(flat_avg, (word_ids.len(), hidden_size), device)?;
  let gathered = matrix.index_select(&ids, 0)?;
  let projected = avg.unsqueeze(1)?.matmul(&gathered)?.squeeze(1)?;
  let normalized = l2_normalize(&projected.tanh()?)?;
  let values = normalized.to_vec2::<f32>()?;
  let mut pairs = Vec::with_capacity(word_ids.len() * index::EMBEDDING_SIZE as usize);
  for (row, &word_id) in values.iter().zip(&word_ids) {
    let tf = if uniform_weight {
      1.0
    } else {
      term_frequency::saturation(groups[&word_id].word_count, doc_len)
    };
    for (component, &value) in row.iter().enumerate() {
      pairs.push((index::dense_index(word_id, component as u32), value * tf));
    }
  }
  Ok(pairs)
}