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

//! Per-row scoring from reconstructed words and pooled attention. Unlike miniCOIL, a word is dropped only as a
//! stopword or one of Python's 32 ASCII punctuation characters, so inner punctuation survives.

use std::collections::HashMap;

use super::super::hash::bm25_hash;
use super::super::stemmer::stem;
use super::super::stopwords::Stopwords;
use super::super::words;
use crate::sparse::SparseEmbedding;

/// `Bm42`'s `alpha` default, which fastembed's catalog keeps for this model.
const ALPHA: f32 = 0.5;

const SPECIAL_TOKENS: [&str; 5] = ["[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"];

fn is_special_token(word: &str) -> bool {
  SPECIAL_TOKENS.contains(&word)
}

/// Python's `string.punctuation`, in the library's order.
const ASCII_PUNCTUATION: &str = "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~";

/// Matches only a whole one-character word, as the reference's set of one-character strings does.
fn is_ascii_punctuation(word: &str) -> bool {
  let mut chars = word.chars();
  match (chars.next(), chars.next()) {
    (Some(only), None) => ASCII_PUNCTUATION.contains(only),
    _ => false,
  }
}

/// `_filter_pair_tokens`' drop, plus the special-token drop `_reconstruct_bpe` applies one step earlier.
fn is_dropped(word: &str, stopwords: &Stopwords) -> bool {
  is_special_token(word) || stopwords.contains(word) || is_ascii_punctuation(word)
}

/// One document row: each surviving word sums `pooled` over its token span, the max per stem is kept, then
/// rescored as `ln(1 + value) ^ alpha` and hashed. Specials are reconstructed too and dropped by name.
pub(super) fn pool_row(
  pooled: &[f32],
  ids: &[u32],
  token_text: &[String],
  stopwords: &Stopwords,
) -> SparseEmbedding {
  let texts: Vec<String> = ids
    .iter()
    .map(|&id| token_text.get(id as usize).cloned().unwrap_or_default())
    .collect();
  let mut max_weight: HashMap<String, f32> = HashMap::new();
  for word in words::reconstruct(&texts) {
    if is_dropped(&word.text, stopwords) {
      continue;
    }
    let weight: f32 = word.tokens.clone().map(|position| pooled[position]).sum();
    let stemmed = stem(&word.text);
    let entry = max_weight.entry(stemmed).or_insert(0.0);
    if weight > *entry {
      *entry = weight;
    }
  }
  SparseEmbedding::from_pairs(
    max_weight
      .into_iter()
      .map(|(stemmed, weight)| (bm25_hash(&stemmed), (1.0 + weight).ln().powf(ALPHA))),
  )
}

/// The query path: every unique surviving stem at a flat weight of `1.0`, no attention.
pub(super) fn pool_query_row(
  ids: &[u32],
  token_text: &[String],
  stopwords: &Stopwords,
) -> SparseEmbedding {
  let texts: Vec<String> = ids
    .iter()
    .map(|&id| token_text.get(id as usize).cloned().unwrap_or_default())
    .collect();
  let pairs = words::reconstruct(&texts)
    .into_iter()
    .filter(|word| !is_dropped(&word.text, stopwords))
    .map(|word| (bm25_hash(&stem(&word.text)), 1.0f32));
  SparseEmbedding::from_pairs(pairs)
}

#[cfg(test)]
mod tests;