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

use candle_core::DType;

use super::DtypeSupport;
use super::Family;
use super::Support;
use super::WeightsFormat;
use super::entry::Entry;
use super::t4::T4_BF16_UNTESTED;
use super::t4::T4_F16_NAN_REFUSED;
use crate::pooling::Pooling;

/// The six's shared fields: mean pooling, 8192 default, no Dense, prompts or Matryoshka, Apache-2.0, ungated.
/// `default_dtype` is F32 although every checkpoint is F16 on disk; `size_gb` is `model.safetensors`' size.
const fn jina(
  id: &'static str,
  dim: usize,
  size_gb: f32,
  backends: Support,
  dtypes: DtypeSupport,
  accel_only_dtypes: &'static [DType],
) -> Entry {
  Entry {
    id,
    family: Family::JinaBert,
    dim,
    default_max_length: 8192,
    pooling: Pooling::Mean,
    cpu: backends,
    metal: backends,
    cuda: Support::Verified,
    query_template: None,
    passage_template: None,
    default_instruction: None,
    dense_dirs: &[],
    mrl: false,
    layer_norm_after_pooling: false,
    gated: false,
    default_dtype: DType::F32,
    dtypes,
    license: "Apache-2.0",
    size_gb,
    phase: 2,
    weights: WeightsFormat::Safetensors,
    accel_only_dtypes,
    cell_exceptions: T4_BF16_UNTESTED,
  }
}

/// `jinaai/jina-embeddings-v2-small-en`: the base block shape, 512 wide with eight heads on the power-of-two ALiBi
/// branch; no `Normalize` module, and at its 8192 default an eight-head additive mask of 2.1 GB at batch 1.
const SMALL_EN: Entry = jina(
  "jinaai/jina-embeddings-v2-small-en",
  512,
  0.07,
  Support::Verified,
  DtypeSupport::verified(),
  &[],
);

/// `jinaai/jina-embeddings-v2-base-en`: the base block shape with twelve heads, whose ALiBi slopes take the
/// interpolated branch and are not monotonic in head order.
const BASE_EN: Entry = jina(
  "jinaai/jina-embeddings-v2-base-en",
  768,
  0.27,
  Support::Verified,
  DtypeSupport::verified(),
  &[],
);

/// `jinaai/jina-embeddings-v2-base-de`: the base block shape over a byte-level BPE that keeps a whitespace-only input
/// as six tokens; it pads with its `pad_token_id`, `<s>`, and its F16 is `NaN` on the CPU and the T4, refused on both.
const BASE_DE: Entry = Entry {
  cell_exceptions: T4_F16_NAN_REFUSED,
  ..jina(
    "jinaai/jina-embeddings-v2-base-de",
    768,
    0.32,
    Support::Verified,
    DtypeSupport::verified(),
    &[DType::F16],
  )
};

/// `jinaai/jina-embeddings-v2-base-zh`: the base block shape over its own NFC and lowercasing tokenizer, which drops
/// a whitespace-only input; its F16 is `NaN` for some inputs on the CPU and the T4, so both refuse it.
const BASE_ZH: Entry = Entry {
  cell_exceptions: T4_F16_NAN_REFUSED,
  ..jina(
    "jinaai/jina-embeddings-v2-base-zh",
    768,
    0.32,
    Support::Verified,
    DtypeSupport::verified(),
    &[DType::F16],
  )
};

/// `jinaai/jina-embeddings-v2-base-es`: the QK-norm block shape, a LayerNorm over the full query and key
/// projections before the heads split; read as the base shape, it loads and returns wrong vectors.
const BASE_ES: Entry = jina(
  "jinaai/jina-embeddings-v2-base-es",
  768,
  0.32,
  Support::Verified,
  DtypeSupport::verified(),
  &[],
);

/// `jinaai/jina-embeddings-v2-base-code`: the split-norm block shape, whose feed-forward gates its second half and
/// has no norm or residual of its own; no `Normalize` module.
const BASE_CODE: Entry = jina(
  "jinaai/jina-embeddings-v2-base-code",
  768,
  0.32,
  Support::Verified,
  DtypeSupport::verified(),
  &[],
);

pub(super) const ENTRIES: &[Entry] = &[SMALL_EN, BASE_EN, BASE_DE, BASE_ZH, BASE_ES, BASE_CODE];