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

//! `Bm42Sparse`: the last BERT layer's CLS-attention row, reduced to one score per whole word.

use candle_core::DType;

use super::pooling;
use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::SparseModel;
use crate::models::bert::BertModel;
use crate::models::sparse::stopwords::Stopwords;
use crate::sparse::SparseEmbedding;

/// `all-MiniLM-L6-v2`'s BM42 head. It holds the concrete `BertModel`, not `Box<dyn Encoder>`, because it
/// needs `forward_with_attention`, which the `Encoder` trait lacks.
pub struct Bm42Sparse {
  bert: BertModel,
  vocab_size: usize,
  token_text: Vec<String>,
  stopwords: Stopwords,
}

impl Bm42Sparse {
  pub fn new(bert: BertModel, vocab_size: usize, token_text: Vec<String>) -> Self {
    Self {
      bert,
      vocab_size,
      token_text,
      stopwords: Stopwords::parse(include_str!("../english_stopwords.txt")),
    }
  }

  /// The last layer's CLS attention averaged over heads; padded keys are already zero, so no mask multiply.
  fn pooled_attention(&self, input: &EncoderInput) -> Result<Vec<Vec<f32>>> {
    let (_hidden, weights) = self.bert.forward_with_attention(input)?;
    let cls_row = weights.narrow(2, 0, 1)?.squeeze(2)?;
    let pooled = cls_row.mean(1)?.to_dtype(DType::F32)?;
    Ok(pooled.to_vec2::<f32>()?)
  }
}

impl SparseModel for Bm42Sparse {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let pooled_rows = self.pooled_attention(input)?;
    let id_rows = input.input_ids.to_vec2::<u32>()?;
    Ok(
      pooled_rows
        .iter()
        .zip(id_rows.iter())
        .zip(input.lengths.iter())
        .map(|((pooled, ids), &len)| {
          let len = len.min(ids.len()).min(pooled.len());
          pooling::pool_row(
            &pooled[..len],
            &ids[..len],
            &self.token_text,
            &self.stopwords,
          )
        })
        .collect(),
    )
  }

  /// The query path never runs the encoder, as in `Bm42.query_embed`: reconstructed words at a flat weight.
  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let id_rows = input.input_ids.to_vec2::<u32>()?;
    Ok(
      id_rows
        .iter()
        .zip(input.lengths.iter())
        .map(|(ids, &len)| {
          let len = len.min(ids.len());
          pooling::pool_query_row(&ids[..len], &self.token_text, &self.stopwords)
        })
        .collect(),
    )
  }

  fn vocab_size(&self) -> usize {
    self.vocab_size
  }

  fn max_position_embeddings(&self) -> usize {
    self.bert.max_position_embeddings()
  }

  fn weights_dtype(&self) -> DType {
    self.bert.weights_dtype()
  }
}

#[cfg(test)]
mod tests;