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

//! Windowed scoring, as the reference does it: every window of every document batched `batch_size` at a time in
//! document order, each document scored by its best window; a window's batch moves its score by at most 5.4e-7.

use tokenizers::Encoding;

use super::TextCrossEncoder;
use crate::error::CanddingError;
use crate::error::Result;
use crate::rerank::RerankScore;
use crate::tokenizer::WindowedTokenizer;

impl TextCrossEncoder {
  pub(super) fn score_windows<S: AsRef<str>>(
    &self,
    tokenizer: &WindowedTokenizer,
    query: &str,
    documents: &[S],
    batch_size: usize,
  ) -> Result<Vec<RerankScore>> {
    let mut owners = Vec::new();
    let mut windows: Vec<Encoding> = Vec::new();
    for (index, document) in documents.iter().enumerate() {
      let cut = tokenizer.windows(query, document.as_ref())?;
      owners.extend(std::iter::repeat_n(index, cut.len()));
      windows.extend(cut);
    }
    let mut scores: Vec<Vec<RerankScore>> = vec![Vec::new(); documents.len()];
    for (owners, windows) in owners.chunks(batch_size).zip(windows.chunks(batch_size)) {
      let batch = self.forward(|device| tokenizer.batch(windows, device))?;
      for (&owner, score) in owners.iter().zip(batch) {
        scores[owner].push(score);
      }
    }
    scores.iter().map(|scores| best_window(scores)).collect()
  }
}

/// A document's score: its best window's, by `RerankScore`'s order. The reference's `merge_scores` keeps the largest
/// of a document's window sigmoids; a document always has a window, the empty one included.
fn best_window(scores: &[RerankScore]) -> Result<RerankScore> {
  scores.iter().copied().max().ok_or_else(|| {
    CanddingError::Tokenizer("a windowed document produced no window to score".to_string())
  })
}