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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197

//! Compact static form of a multi-vector catalog entry and its public runtime form, parallel to `sparse::entry`.

use candle_core::DType;
use serde::Deserialize;
use serde::Serialize;

use crate::models::MultiVectorHead;
use crate::models::xlm_roberta_flash::PositionScheme;
use crate::multi_vector::MaskedRowFill;
use crate::multi_vector::ScoreReduction;
use crate::registry::BackendSupport;
use crate::registry::CellException;
use crate::registry::DtypeBackendSupport;
use crate::registry::DtypeSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;
use crate::registry::multi_vector::conventions::HeadConventions;

/// Compact static form of a multi-vector catalog entry. It names no second repository, since every checkpoint keeps
/// its projection beside its encoder, and has no `vocab_size` or pooling: the projection fields below are its shape.
pub(crate) struct MultiVectorEntry {
  pub id: &'static str,
  pub family: Family,
  pub head: MultiVectorHead,
  pub default_max_length: usize,
  /// Whether the classification row survives the head: `BAAI/bge-m3` drops it, the ColBERT-style checkpoints keep it.
  pub keeps_cls_row: bool,
  /// The encoder width the projection reads; not derivable from `family`, and checked against the encoder at load.
  pub projection_in: usize,
  /// The width of every returned row: 1024 for bge-m3, 128 for colbertv2, 96 for answerai-colbert-small.
  pub projection_out: usize,
  /// Whether the projection carries a bias: `true` for bge-m3's `colbert_linear.pt`, `false` for the ColBERT roots.
  pub projection_bias: bool,
  /// Whether the head L2-normalizes every row after the projection; `true` for every entry.
  pub normalize: bool,
  /// How the checkpoint's reference reduces the MaxSim sum: divided by the query length for `BAAI/bge-m3`, plain
  /// otherwise.
  pub score_reduction: ScoreReduction,
  /// The prefix a query gets before tokenizing, registered as a special token so it stays one piece; `None` for
  /// `BAAI/bge-m3`.
  pub query_marker: Option<&'static str>,
  /// The prefix a document gets before tokenizing, registered as a special token; `None` for bge-m3.
  pub document_marker: Option<&'static str>,
  /// The fixed row count a query is padded or truncated to with the mask token, a contract of the checkpoint rather
  /// than a `max_length`; `None` for `BAAI/bge-m3`.
  pub query_length: Option<usize>,
  /// Whether the encoder's self-attention sees a query's mask-token padding: `true` for `jinaai/jina-colbert-v2` alone.
  pub attends_to_mask_tokens: bool,
  /// The value a masked document row takes before the per-query-row maximum: `Zero` for bge-m3 (FlagEmbedding),
  /// `DtypeMinimum` for the sentence-transformers checkpoints.
  pub masked_document_row_fill: MaskedRowFill,
  /// How a flash XLM-RoBERTa checkpoint numbers its tokens (`Some(Rotary)` for jina-colbert-v2, else `None`);
  /// the encoder's load refuses a `config.json` naming another scheme.
  pub position_scheme: Option<PositionScheme>,
  pub cpu: Support,
  pub metal: Support,
  pub cuda: Support,
  pub dtypes: DtypeSupport,
  pub license: &'static str,
  pub size_gb: f32,
  pub phase: u8,
  pub weights: WeightsFormat,
  pub gated: bool,
  pub accel_only_dtypes: &'static [DType],
  /// Cells a run decided that the derivation cannot state, written over its rows (`CellException`).
  pub cell_exceptions: &'static [CellException],
}

impl MultiVectorEntry {
  pub fn descriptor(&self) -> MultiVectorDescriptor {
    let backends = BackendSupport {
      cpu: self.cpu,
      metal: self.metal,
      cuda: self.cuda,
    };
    MultiVectorDescriptor {
      id: self.id.to_string(),
      family: self.family,
      head: self.head,
      default_max_length: self.default_max_length,
      keeps_cls_row: self.keeps_cls_row,
      projection_in: self.projection_in,
      projection_out: self.projection_out,
      projection_bias: self.projection_bias,
      normalize: self.normalize,
      score_reduction: self.score_reduction,
      query_marker: self.query_marker.map(str::to_string),
      document_marker: self.document_marker.map(str::to_string),
      query_length: self.query_length,
      attends_to_mask_tokens: self.attends_to_mask_tokens,
      masked_document_row_fill: self.masked_document_row_fill,
      position_scheme: self.position_scheme,
      backends,
      dtype_backends: DtypeBackendSupport::derive(backends, self.dtypes, self.accel_only_dtypes)
        .with_exceptions(self.cell_exceptions),
      license: self.license.to_string(),
      size_gb: self.size_gb,
      phase: self.phase,
      weights: self.weights,
      gated: self.gated,
    }
  }

  /// This entry as `registry::support::rules` reads it.
  #[cfg(test)]
  pub(crate) fn cells(&self) -> crate::registry::support::rules::Cells<'static> {
    let descriptor = self.descriptor();
    crate::registry::support::rules::Cells {
      id: self.id,
      backends: descriptor.backends,
      dtypes: self.dtypes,
      accel_only: self.accel_only_dtypes,
      exceptions: self.cell_exceptions,
      rows: descriptor.dtype_backends,
    }
  }
}

/// Everything the library knows about a multi-vector model beyond its weights; `new` builds one for a head outside
/// the registry.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct MultiVectorDescriptor {
  pub id: String,
  pub family: Family,
  pub head: MultiVectorHead,
  pub default_max_length: usize,
  /// Whether the classification position survives as a row; `false` for `BAAI/bge-m3`.
  pub keeps_cls_row: bool,
  pub projection_in: usize,
  pub projection_out: usize,
  pub projection_bias: bool,
  /// Whether every row is L2-normalized; the builder's `normalize` overrides it.
  pub normalize: bool,
  /// How `late_interaction_score_for` reduces the MaxSim sum.
  pub score_reduction: ScoreReduction,
  /// The prefix a query gets before tokenizing, such as `[unused0] `.
  pub query_marker: Option<String>,
  /// The prefix a document gets before tokenizing, such as `[unused1] `.
  pub document_marker: Option<String>,
  /// The fixed row count a query is padded to with the mask token; `None` when there is none.
  pub query_length: Option<usize>,
  /// Whether a query's mask-token padding takes part in attention, as for `jinaai/jina-colbert-v2`.
  pub attends_to_mask_tokens: bool,
  /// The value a masked document row takes before the maximum.
  pub masked_document_row_fill: MaskedRowFill,
  /// How a flash XLM-RoBERTa encoder numbers tokens, `rotary` or `absolute`; `null` for every other family.
  pub position_scheme: Option<PositionScheme>,
  pub backends: BackendSupport,
  /// Per-dtype support broken down by device, the same shape `ModelDescriptor::dtype_backends` carries.
  pub dtype_backends: DtypeBackendSupport,
  pub license: String,
  pub size_gb: f32,
  pub phase: u8,
  pub weights: WeightsFormat,
  pub gated: bool,
}

impl MultiVectorDescriptor {
  /// The facts no default can stand for; the rest comes from `head` (a flash encoder takes rotary positions),
  /// unit-length rows, and `ModelDescriptor::new`'s defaults.
  pub fn new(
    id: impl Into<String>,
    family: Family,
    head: MultiVectorHead,
    projection_in: usize,
    projection_out: usize,
  ) -> Self {
    let conventions = HeadConventions::of(head);
    Self {
      id: id.into(),
      family,
      head,
      default_max_length: 512,
      keeps_cls_row: conventions.keeps_cls_row,
      projection_in,
      projection_out,
      projection_bias: conventions.projection_bias,
      normalize: true,
      score_reduction: conventions.score_reduction,
      query_marker: conventions.query_marker.map(str::to_string),
      document_marker: conventions.document_marker.map(str::to_string),
      query_length: conventions.query_length,
      attends_to_mask_tokens: conventions.attends_to_mask_tokens,
      masked_document_row_fill: conventions.masked_document_row_fill,
      position_scheme: (family == Family::XlmRobertaFlash).then_some(PositionScheme::Rotary),
      backends: BackendSupport::untested(),
      dtype_backends: DtypeBackendSupport::untested(),
      license: "unknown".to_string(),
      size_gb: 0.0,
      phase: u8::MAX,
      weights: conventions.weights,
      gated: false,
    }
  }
}