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

//! Which positions a flash XLM-RoBERTa checkpoint reads: stated by its registry entry and checked against
//! `config.json`, since the two implementation repositories default `position_embedding_type` differently.

use serde::Deserialize;
use serde::Serialize;

use crate::error::CanddingError;
use crate::error::Result;

/// How a flash checkpoint numbers its tokens, spelled as `config.json`'s `position_embedding_type` spells it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum PositionScheme {
  /// Rotary angles on the query and the key inside attention, at positions `0..seq`, with no position table:
  /// `jinaai/jina-colbert-v2`.
  Rotary,
  /// A learned table added to the token embeddings, numbered from `pad_token_id + 1`:
  /// `jinaai/jina-reranker-v2-base-multilingual`.
  Absolute,
}

impl PositionScheme {
  /// The serde name, the `position_embedding_type` value that states this scheme.
  pub fn name(self) -> &'static str {
    match self {
      PositionScheme::Rotary => "rotary",
      PositionScheme::Absolute => "absolute",
    }
  }

  /// The scheme a `position_embedding_type` value names; `None` for one this family lacks, such as `alibi`.
  #[doc(hidden)]
  pub fn from_config(value: &str) -> Option<Self> {
    [PositionScheme::Rotary, PositionScheme::Absolute]
      .into_iter()
      .find(|scheme| scheme.name() == value)
  }

  /// The scheme a load runs: `declared` when `config.json` names the same one.
  /// A missing declaration, a missing config value or a disagreement is refused.
  #[doc(hidden)]
  pub fn resolve(declared: Option<Self>, config: Option<&str>) -> Result<Self> {
    match declared {
      Some(scheme) if config == Some(scheme.name()) => Ok(scheme),
      _ => Err(CanddingError::PositionSchemeDisagreesWithConfig {
        declared: declared.map_or("none", PositionScheme::name).to_string(),
        config: config.unwrap_or("none").to_string(),
      }),
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  /// The serde spelling is the `config.json` one, which `resolve` compares against and every surface prints.
  #[test]
  fn each_scheme_serializes_as_the_config_value_that_names_it() {
    for scheme in [PositionScheme::Rotary, PositionScheme::Absolute] {
      let json = serde_json::to_value(scheme).unwrap();
      assert_eq!(json, serde_json::Value::String(scheme.name().to_string()));
      assert_eq!(PositionScheme::from_config(scheme.name()), Some(scheme));
    }
    assert_eq!(PositionScheme::from_config("alibi"), None);
    assert_eq!(PositionScheme::from_config("relative_key"), None);
  }

  #[test]
  fn resolve_keeps_the_declared_scheme_when_the_config_names_it() {
    for scheme in [PositionScheme::Rotary, PositionScheme::Absolute] {
      assert_eq!(
        PositionScheme::resolve(Some(scheme), Some(scheme.name())).unwrap(),
        scheme
      );
    }
  }

  #[test]
  fn resolve_refuses_a_disagreement_and_either_side_missing() {
    for (declared, config, expected) in [
      (
        Some(PositionScheme::Absolute),
        Some("rotary"),
        ("absolute", "rotary"),
      ),
      (
        Some(PositionScheme::Rotary),
        Some("absolute"),
        ("rotary", "absolute"),
      ),
      (None, Some("absolute"), ("none", "absolute")),
      (Some(PositionScheme::Absolute), None, ("absolute", "none")),
      (
        Some(PositionScheme::Rotary),
        Some("alibi"),
        ("rotary", "alibi"),
      ),
    ] {
      match PositionScheme::resolve(declared, config) {
        Err(CanddingError::PositionSchemeDisagreesWithConfig { declared, config }) => {
          assert_eq!((declared.as_str(), config.as_str()), expected)
        }
        other => panic!("{declared:?} against {config:?}: {other:?}"),
      }
    }
  }
}