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

//! Per-row pooling: words are rebuilt and averaged in host memory first, so the only gather selects one `(512, 4)`
//! slice per distinct word in the row (64 MiB at most), never a `(batch, seq, 512, 4)` tensor (512 MiB at 8 × 8192).

mod clean_words;
mod gather;

use std::collections::HashMap;

use candle_core::Tensor;

use super::index;
use super::term_frequency;
use super::vocab::Vocab;
use crate::error::Result;
use crate::models::sparse::punctuation::is_bare_punctuation;
use crate::models::sparse::stemmer::stem;
use crate::models::sparse::stopwords::Stopwords;
use crate::models::sparse::words;
use crate::sparse::SparseEmbedding;

/// The five BERT specials, removed as words with punctuation and stopwords as the reference's `clean_words` does,
/// so the `[CLS]` and `[SEP]` every row carries never reach the unrecognized-word hash.
const SPECIAL_TOKENS: [&str; 5] = ["[CLS]", "[SEP]", "[PAD]", "[UNK]", "[MASK]"];

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

/// Unrecognized pieces group by stem, since `hashed_index` hashes the stem; `representative` keeps one unstemmed
/// survivor because stemming twice is not assumed to be a no-op.
struct UnrecognizedGroup {
  representative: String,
  word_count: u32,
}

/// One `(hash index, tf weight)` pair per distinct unrecognized stem; these never reach the per-word matrix.
fn unrecognized_pairs(
  groups: &HashMap<String, UnrecognizedGroup>,
  doc_len: u32,
  uniform_weight: bool,
) -> Vec<(u32, f32)> {
  groups
    .values()
    .map(|group| {
      let tf = if uniform_weight {
        1.0
      } else {
        term_frequency::saturation(group.word_count, doc_len)
      };
      (index::hashed_index(&group.representative), tf)
    })
    .collect()
}

/// One row: rebuild words, drop specials and stopwords, gather and normalize known words, hash the rest.
/// `uniform_weight` sets every term weight to 1.0 for queries; negative components stay, exact zeros drop.
pub(super) fn pool_row(
  hidden_row: &[Vec<f32>],
  ids_row: &[u32],
  token_text: &[String],
  vocab: &Vocab,
  stopwords: &Stopwords,
  matrix: &Tensor,
  uniform_weight: bool,
) -> Result<SparseEmbedding> {
  let hidden_size = matrix.dims()[1];
  let texts: Vec<String> = ids_row
    .iter()
    .map(|&id| token_text.get(id as usize).cloned().unwrap_or_default())
    .collect();
  let mut recognized: HashMap<u32, gather::RecognizedGroup> = HashMap::new();
  let mut unrecognized: HashMap<String, UnrecognizedGroup> = HashMap::new();
  let mut doc_len: u32 = 0;
  for word in words::reconstruct(&texts) {
    if is_special_token(&word.text)
      || stopwords.contains(&word.text)
      || is_bare_punctuation(&word.text)
    {
      continue;
    }
    match vocab.resolve(&word.text) {
      Some(word_id) => {
        doc_len += 1;
        let group = recognized
          .entry(word_id)
          .or_insert_with(|| gather::RecognizedGroup::new(hidden_size));
        for position in word.tokens {
          group.add_token(&hidden_row[position]);
        }
        group.word_count += 1;
      }
      None => {
        for subword in clean_words::resplit(&word.text, stopwords) {
          doc_len += 1;
          let stem = stem(&subword);
          unrecognized
            .entry(stem)
            .or_insert_with(move || UnrecognizedGroup {
              representative: subword,
              word_count: 0,
            })
            .word_count += 1;
        }
      }
    }
  }
  let mut pairs = gather::pairs(&recognized, matrix, doc_len, uniform_weight)?;
  pairs.extend(unrecognized_pairs(&unrecognized, doc_len, uniform_weight));
  Ok(SparseEmbedding::from_pairs(pairs))
}