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

//! `flash_fixture`'s position table moves the logit by more than 0.05 when skipped or numbered from 0.

use super::*;
use crate::error::Result;
use crate::models::rerank::flash_fixture;
use crate::models::rerank::flash_fixture::Numbering;

fn flash(
  score_activation: ScoreActivation,
  declared: Option<PositionScheme>,
) -> Result<RobertaClassificationRerank> {
  let vb = VarBuilder::from_tensors(flash_fixture::tensors(), DType::F32, &Device::Cpu);
  RobertaClassificationRerank::load(
    xlm_roberta_flash::ARCHITECTURE,
    &flash_fixture::config_json("absolute"),
    vb,
    score_activation,
    declared,
  )
}

/// Two rows starting with `<s>`, id 0, the second right-padded with the pad id: the XLM-RoBERTa pair layout.
fn flash_scores(model: &RobertaClassificationRerank) -> Vec<f32> {
  let input_ids = Tensor::new(
    &[[0u32, 5, 2, 2, 7, 2], [0u32, 9, 2, 2, 2, 1]],
    &Device::Cpu,
  )
  .unwrap();
  let attention_mask = Tensor::new(
    &[[1u32, 1, 1, 1, 1, 1], [1u32, 1, 1, 1, 1, 0]],
    &Device::Cpu,
  )
  .unwrap();
  let token_type_ids = input_ids.zeros_like().unwrap();
  let lengths = vec![6, 5];
  let input = EncoderInput {
    input_ids: &input_ids,
    attention_mask: &attention_mask,
    token_type_ids: &token_type_ids,
    lengths: &lengths,
  };
  model
    .forward(&input)
    .unwrap()
    .into_iter()
    .map(|s| s.value())
    .collect()
}

#[test]
fn the_flash_head_reads_the_learned_table_numbered_from_pad_plus_one() {
  let model = flash(ScoreActivation::Identity, Some(PositionScheme::Absolute)).unwrap();
  let expected = flash_fixture::expected_logit(0, Numbering::FromPadPlusOne);
  let scores = flash_scores(&model);
  assert_eq!(scores.len(), 2);
  for score in scores {
    assert!(
      (score as f64 - expected).abs() < 1e-5,
      "{score} against {expected}, the table numbered from 0 gives {}, no table {}",
      flash_fixture::expected_logit(0, Numbering::FromZero),
      flash_fixture::expected_logit(0, Numbering::NoTable)
    );
  }
}

#[test]
fn the_flash_head_applies_its_activation() {
  let model = flash(ScoreActivation::Sigmoid, Some(PositionScheme::Absolute)).unwrap();
  let expected = ScoreActivation::Sigmoid
    .apply(flash_fixture::expected_logit(0, Numbering::FromPadPlusOne) as f32);
  for score in flash_scores(&model) {
    assert!(
      (score - expected).abs() < 1e-6,
      "{score} against {expected}"
    );
  }
}

#[test]
fn the_flash_head_hands_its_encoder_the_scheme_it_was_given() {
  for declared in [Some(PositionScheme::Rotary), None] {
    match flash(ScoreActivation::Sigmoid, declared) {
      Err(CanddingError::PositionSchemeDisagreesWithConfig { config, .. }) => {
        assert_eq!(config, "absolute")
      }
      Ok(_) => panic!("{declared:?} against an absolute config must be refused"),
      Err(other) => panic!("{declared:?}: unexpected error: {other}"),
    }
  }
}