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

//! `Qdrant/bm25`'s document and query rows: drop punctuation, stopwords and tokens over 40 characters, stem,
//! weight and hash, in the reference's `_stem` order; a hash collision keeps the larger weight.

use std::collections::HashMap;
use std::collections::HashSet;

use super::super::hash::bm25_hash;
use super::super::punctuation::is_bare_punctuation;
use super::super::stemmer::stem;
use super::super::stopwords::Stopwords;
use super::term_frequency;
use super::tokenizer;
use crate::sparse::SparseEmbedding;

/// `Bm25.__init__`'s `token_max_length` default, which fastembed's catalog keeps for this model.
const TOKEN_MAX_LENGTH: usize = 40;

/// `_stem`; `tokens` is already lowercase, so the reference's `token.lower()` is implicit.
fn stem_tokens(tokens: Vec<String>, stopwords: &Stopwords) -> Vec<String> {
  tokens
    .into_iter()
    .filter(|token| !is_bare_punctuation(token))
    .filter(|token| !stopwords.contains(token))
    .filter(|token| token.chars().count() <= TOKEN_MAX_LENGTH)
    .map(|token| stem(&token))
    .filter(|stemmed| !stemmed.is_empty())
    .collect()
}

fn tokenize_and_stem(text: &str, stopwords: &Stopwords) -> Vec<String> {
  let cleaned = tokenizer::remove_non_alphanumeric(text);
  stem_tokens(tokenizer::tokenize(&cleaned), stopwords)
}

/// `Bm25.raw_embed`: count stems, weight each by `term_frequency::saturation`, then hash.
pub(crate) fn embed_document(text: &str, stopwords: &Stopwords) -> SparseEmbedding {
  let stems = tokenize_and_stem(text, stopwords);
  let doc_len = stems.len() as u32;
  let mut counts: HashMap<String, u32> = HashMap::new();
  for stemmed in stems {
    *counts.entry(stemmed).or_insert(0) += 1;
  }
  SparseEmbedding::from_pairs(counts.into_iter().map(|(stemmed, count)| {
    (
      bm25_hash(&stemmed),
      term_frequency::saturation(count, doc_len),
    )
  }))
}

/// `Bm25.query_embed`: every unique surviving stem at a flat weight of `1.0`.
pub(crate) fn embed_query(text: &str, stopwords: &Stopwords) -> SparseEmbedding {
  let stems: HashSet<String> = tokenize_and_stem(text, stopwords).into_iter().collect();
  SparseEmbedding::from_pairs(
    stems
      .into_iter()
      .map(|stemmed| (bm25_hash(&stemmed), 1.0f32)),
  )
}

/// The published English stopword list the builder compiles in, byte-identical to `Qdrant/bm25`'s `english.txt`.
#[cfg(test)]
fn test_stopwords() -> Stopwords {
  Stopwords::parse(include_str!("../english_stopwords.txt"))
}

#[cfg(test)]
mod tests;