llm_kernel/embedding/
catalog.rs

1//! Zero-dep embedding model catalog.
2//!
3//! Mirrors `fastembed::EmbeddingModel` (44 variants) so the catalog is always
4//! available — even when the `embedding-fastembed` feature is disabled.
5
6/// Embedding model catalog with metadata for all supported ONNX models.
7///
8/// Variant names match `fastembed::EmbeddingModel` exactly for trivial 1:1 mapping.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
10pub enum EmbeddingModel {
11    // ── sentence-transformers ───────────────────────
12    /// BGE Small EN v1.5 — fast 384-dim English model (default).
13    #[default]
14    BGESmallENV15,
15    /// sentence-transformers all-MiniLM-L6-v2 (384-dim).
16    AllMiniLML6V2,
17    /// Quantized all-MiniLM-L6-v2 (384-dim).
18    AllMiniLML6V2Q,
19    /// sentence-transformers all-MiniLM-L12-v2 (384-dim).
20    AllMiniLML12V2,
21    /// Quantized all-MiniLM-L12-v2 (384-dim).
22    AllMiniLML12V2Q,
23    /// sentence-transformers all-mpnet-base-v2 (768-dim).
24    AllMpnetBaseV2,
25    // ── BGE family ──────────────────────────────────
26    /// BGE Base EN v1.5 (768-dim).
27    BGEBaseENV15,
28    /// Quantized BGE Base EN v1.5 (768-dim).
29    BGEBaseENV15Q,
30    /// BGE Large EN v1.5 (1024-dim).
31    BGELargeENV15,
32    /// Quantized BGE Large EN v1.5 (1024-dim).
33    BGELargeENV15Q,
34    /// Quantized BGE Small EN v1.5 (384-dim).
35    BGESmallENV15Q,
36    /// BGE Small ZH v1.5 — Chinese model (512-dim).
37    BGESmallZHV15,
38    /// BGE Large ZH v1.5 — Chinese model (1024-dim).
39    BGELargeZHV15,
40    /// BGE M3 — multilingual, 100+ languages, 8192 context (1024-dim).
41    BGEM3,
42    // ── Nomic ───────────────────────────────────────
43    /// Nomic Embed Text v1 — 8192 context English model (768-dim).
44    NomicEmbedTextV1,
45    /// Nomic Embed Text v1.5 — 8192 context English model (768-dim).
46    NomicEmbedTextV15,
47    /// Quantized Nomic Embed Text v1.5 (768-dim).
48    NomicEmbedTextV15Q,
49    // ── Paraphrase ──────────────────────────────────
50    /// Paraphrase multilingual MiniLM-L12-v2 (384-dim).
51    ParaphraseMLMiniLML12V2,
52    /// Quantized paraphrase multilingual MiniLM-L12-v2 (384-dim).
53    ParaphraseMLMiniLML12V2Q,
54    /// Paraphrase multilingual mpnet-base-v2 (768-dim).
55    ParaphraseMLMpnetBaseV2,
56    // ── ModernBERT ──────────────────────────────────
57    /// ModernBERT Embed Large (1024-dim).
58    ModernBertEmbedLarge,
59    // ── E5 multilingual ─────────────────────────────
60    /// Multilingual E5 Small (384-dim).
61    MultilingualE5Small,
62    /// Multilingual E5 Base (768-dim).
63    MultilingualE5Base,
64    /// Multilingual E5 Large (1024-dim).
65    MultilingualE5Large,
66    // ── Mixedbread ──────────────────────────────────
67    /// Mixedbread mxbai-embed-large-v1 (1024-dim).
68    MxbaiEmbedLargeV1,
69    /// Quantized mxbai-embed-large-v1 (1024-dim).
70    MxbaiEmbedLargeV1Q,
71    // ── GTE (Alibaba) ──────────────────────────────
72    /// Alibaba GTE Base EN v1.5 (768-dim).
73    GTEBaseENV15,
74    /// Quantized Alibaba GTE Base EN v1.5 (768-dim).
75    GTEBaseENV15Q,
76    /// Alibaba GTE Large EN v1.5 (1024-dim).
77    GTELargeENV15,
78    /// Quantized Alibaba GTE Large EN v1.5 (1024-dim).
79    GTELargeENV15Q,
80    // ── CLIP ────────────────────────────────────────
81    /// OpenAI CLIP ViT-B/32 text encoder (512-dim).
82    ClipVitB32,
83    // ── Jina ────────────────────────────────────────
84    /// Jina Embeddings v2 Base for code (768-dim).
85    JinaEmbeddingsV2BaseCode,
86    /// Jina Embeddings v2 Base for English (768-dim).
87    JinaEmbeddingsV2BaseEN,
88    // ── Gemma ───────────────────────────────────────
89    /// Google EmbeddingGemma 300M (768-dim).
90    EmbeddingGemma300M,
91    // ── Snowflake Arctic ────────────────────────────
92    /// Snowflake Arctic Embed XS (384-dim).
93    SnowflakeArcticEmbedXS,
94    /// Quantized Snowflake Arctic Embed XS (384-dim).
95    SnowflakeArcticEmbedXSQ,
96    /// Snowflake Arctic Embed S (384-dim).
97    SnowflakeArcticEmbedS,
98    /// Quantized Snowflake Arctic Embed S (384-dim).
99    SnowflakeArcticEmbedSQ,
100    /// Snowflake Arctic Embed M (768-dim).
101    SnowflakeArcticEmbedM,
102    /// Quantized Snowflake Arctic Embed M (768-dim).
103    SnowflakeArcticEmbedMQ,
104    /// Snowflake Arctic Embed M Long — 2048 context (768-dim).
105    SnowflakeArcticEmbedMLong,
106    /// Quantized Snowflake Arctic Embed M Long (768-dim).
107    SnowflakeArcticEmbedMLongQ,
108    /// Snowflake Arctic Embed L (1024-dim).
109    SnowflakeArcticEmbedL,
110    /// Quantized Snowflake Arctic Embed L (1024-dim).
111    SnowflakeArcticEmbedLQ,
112}
113
114impl EmbeddingModel {
115    /// Embedding dimensionality.
116    pub const fn dimension(self) -> usize {
117        match self {
118            // 384-dim
119            Self::BGESmallENV15
120            | Self::AllMiniLML6V2
121            | Self::AllMiniLML6V2Q
122            | Self::AllMiniLML12V2
123            | Self::AllMiniLML12V2Q
124            | Self::BGESmallENV15Q
125            | Self::ParaphraseMLMiniLML12V2
126            | Self::ParaphraseMLMiniLML12V2Q
127            | Self::MultilingualE5Small
128            | Self::SnowflakeArcticEmbedXS
129            | Self::SnowflakeArcticEmbedXSQ
130            | Self::SnowflakeArcticEmbedS
131            | Self::SnowflakeArcticEmbedSQ => 384,
132            // 512-dim
133            Self::BGESmallZHV15 | Self::ClipVitB32 => 512,
134            // 768-dim
135            Self::AllMpnetBaseV2
136            | Self::BGEBaseENV15
137            | Self::BGEBaseENV15Q
138            | Self::NomicEmbedTextV1
139            | Self::NomicEmbedTextV15
140            | Self::NomicEmbedTextV15Q
141            | Self::ParaphraseMLMpnetBaseV2
142            | Self::MultilingualE5Base
143            | Self::GTEBaseENV15
144            | Self::GTEBaseENV15Q
145            | Self::JinaEmbeddingsV2BaseCode
146            | Self::JinaEmbeddingsV2BaseEN
147            | Self::EmbeddingGemma300M
148            | Self::SnowflakeArcticEmbedM
149            | Self::SnowflakeArcticEmbedMQ
150            | Self::SnowflakeArcticEmbedMLong
151            | Self::SnowflakeArcticEmbedMLongQ => 768,
152            // 1024-dim
153            Self::BGELargeENV15
154            | Self::BGELargeENV15Q
155            | Self::BGELargeZHV15
156            | Self::BGEM3
157            | Self::ModernBertEmbedLarge
158            | Self::MultilingualE5Large
159            | Self::MxbaiEmbedLargeV1
160            | Self::MxbaiEmbedLargeV1Q
161            | Self::GTELargeENV15
162            | Self::GTELargeENV15Q
163            | Self::SnowflakeArcticEmbedL
164            | Self::SnowflakeArcticEmbedLQ => 1024,
165        }
166    }
167
168    /// Short human-readable description.
169    pub const fn description(self) -> &'static str {
170        match self {
171            Self::BGESmallENV15 => "v1.5 release of the fast and default English model",
172            Self::AllMiniLML6V2 => "Sentence Transformer model, MiniLM-L6-v2",
173            Self::AllMiniLML6V2Q => "Quantized Sentence Transformer model, MiniLM-L6-v2",
174            Self::AllMiniLML12V2 => "Sentence Transformer model, MiniLM-L12-v2",
175            Self::AllMiniLML12V2Q => "Quantized Sentence Transformer model, MiniLM-L12-v2",
176            Self::AllMpnetBaseV2 => "Sentence Transformer model, mpnet-base-v2",
177            Self::BGEBaseENV15 => "v1.5 release of the base English model",
178            Self::BGEBaseENV15Q => "Quantized v1.5 release of the base English model",
179            Self::BGELargeENV15 => "v1.5 release of the large English model",
180            Self::BGELargeENV15Q => "Quantized v1.5 release of the large English model",
181            Self::BGESmallENV15Q => "Quantized v1.5 release of the fast and default English model",
182            Self::NomicEmbedTextV1 => "8192 context length english model",
183            Self::NomicEmbedTextV15 => "v1.5 release of the 8192 context length english model",
184            Self::NomicEmbedTextV15Q => {
185                "Quantized v1.5 release of the 8192 context length english model"
186            }
187            Self::ParaphraseMLMiniLML12V2 => "Multi-lingual model",
188            Self::ParaphraseMLMiniLML12V2Q => "Quantized multi-lingual model",
189            Self::ParaphraseMLMpnetBaseV2 => {
190                "Sentence-transformers model for clustering or semantic search"
191            }
192            Self::BGESmallZHV15 => "v1.5 release of the small Chinese model",
193            Self::BGELargeZHV15 => "v1.5 release of the large Chinese model",
194            Self::BGEM3 => "Multilingual M3 model with 8192 context length, 100+ languages",
195            Self::ModernBertEmbedLarge => "Large model of ModernBert Text Embeddings",
196            Self::MultilingualE5Small => "Small model of multilingual E5 Text Embeddings",
197            Self::MultilingualE5Base => "Base model of multilingual E5 Text Embeddings",
198            Self::MultilingualE5Large => "Large model of multilingual E5 Text Embeddings",
199            Self::MxbaiEmbedLargeV1 => "Large English embedding model from MixedBreed.ai",
200            Self::MxbaiEmbedLargeV1Q => {
201                "Quantized large English embedding model from MixedBreed.ai"
202            }
203            Self::GTEBaseENV15 => "Base multilingual embedding model from Alibaba",
204            Self::GTEBaseENV15Q => "Quantized base multilingual embedding model from Alibaba",
205            Self::GTELargeENV15 => "Large multilingual embedding model from Alibaba",
206            Self::GTELargeENV15Q => "Quantized large multilingual embedding model from Alibaba",
207            Self::ClipVitB32 => "CLIP text encoder based on ViT-B/32",
208            Self::JinaEmbeddingsV2BaseCode => "Jina embeddings v2 base code",
209            Self::JinaEmbeddingsV2BaseEN => "Jina embeddings v2 base English",
210            Self::EmbeddingGemma300M => "EmbeddingGemma 300M parameter model from Google",
211            Self::SnowflakeArcticEmbedXS => "Snowflake Arctic embed model, xs",
212            Self::SnowflakeArcticEmbedXSQ => "Quantized Snowflake Arctic embed model, xs",
213            Self::SnowflakeArcticEmbedS => "Snowflake Arctic embed model, small",
214            Self::SnowflakeArcticEmbedSQ => "Quantized Snowflake Arctic embed model, small",
215            Self::SnowflakeArcticEmbedM => "Snowflake Arctic embed model, medium",
216            Self::SnowflakeArcticEmbedMQ => "Quantized Snowflake Arctic embed model, medium",
217            Self::SnowflakeArcticEmbedMLong => {
218                "Snowflake Arctic embed model, medium with 2048 context"
219            }
220            Self::SnowflakeArcticEmbedMLongQ => {
221                "Quantized Snowflake Arctic embed model, medium with 2048 context"
222            }
223            Self::SnowflakeArcticEmbedL => "Snowflake Arctic embed model, large",
224            Self::SnowflakeArcticEmbedLQ => "Quantized Snowflake Arctic embed model, large",
225        }
226    }
227
228    /// Optional prefix prepended to query texts before embedding.
229    pub const fn query_prefix(self) -> Option<&'static str> {
230        match self {
231            Self::MultilingualE5Small | Self::MultilingualE5Base | Self::MultilingualE5Large => {
232                Some("query: ")
233            }
234            Self::NomicEmbedTextV15 | Self::NomicEmbedTextV15Q => Some("search_query: "),
235            Self::SnowflakeArcticEmbedXS
236            | Self::SnowflakeArcticEmbedXSQ
237            | Self::SnowflakeArcticEmbedS
238            | Self::SnowflakeArcticEmbedSQ
239            | Self::SnowflakeArcticEmbedM
240            | Self::SnowflakeArcticEmbedMQ
241            | Self::SnowflakeArcticEmbedMLong
242            | Self::SnowflakeArcticEmbedMLongQ
243            | Self::SnowflakeArcticEmbedL
244            | Self::SnowflakeArcticEmbedLQ => {
245                Some("Represent this sentence for searching relevant passages: ")
246            }
247            _ => None,
248        }
249    }
250
251    /// Optional prefix prepended to document texts before embedding.
252    pub const fn doc_prefix(self) -> Option<&'static str> {
253        match self {
254            Self::MultilingualE5Small | Self::MultilingualE5Base | Self::MultilingualE5Large => {
255                Some("passage: ")
256            }
257            Self::NomicEmbedTextV15 | Self::NomicEmbedTextV15Q => Some("search_document: "),
258            _ => None,
259        }
260    }
261
262    /// Approximate ONNX model size in MB.
263    pub const fn size_mb(self) -> usize {
264        match self {
265            // 40 MB
266            Self::BGESmallENV15 | Self::BGESmallENV15Q => 40,
267            // 80 MB
268            Self::AllMiniLML6V2 | Self::AllMiniLML6V2Q => 80,
269            // 90 MB
270            Self::SnowflakeArcticEmbedXS | Self::SnowflakeArcticEmbedXSQ | Self::BGESmallZHV15 => {
271                90
272            }
273            // 120 MB
274            Self::AllMiniLML12V2 | Self::AllMiniLML12V2Q => 120,
275            // 130 MB
276            Self::SnowflakeArcticEmbedS | Self::SnowflakeArcticEmbedSQ => 130,
277            // 260 MB
278            Self::JinaEmbeddingsV2BaseCode
279            | Self::JinaEmbeddingsV2BaseEN
280            | Self::EmbeddingGemma300M => 260,
281            // 420 MB
282            Self::AllMpnetBaseV2
283            | Self::BGEBaseENV15
284            | Self::BGEBaseENV15Q
285            | Self::GTEBaseENV15
286            | Self::GTEBaseENV15Q => 420,
287            // 430 MB
288            Self::SnowflakeArcticEmbedM
289            | Self::SnowflakeArcticEmbedMQ
290            | Self::SnowflakeArcticEmbedMLong
291            | Self::SnowflakeArcticEmbedMLongQ => 430,
292            // 470 MB
293            Self::ParaphraseMLMiniLML12V2
294            | Self::ParaphraseMLMiniLML12V2Q
295            | Self::MultilingualE5Small => 470,
296            // 550 MB
297            Self::NomicEmbedTextV1 | Self::NomicEmbedTextV15 | Self::NomicEmbedTextV15Q => 550,
298            // 600 MB
299            Self::BGEM3 | Self::ModernBertEmbedLarge | Self::ClipVitB32 => 600,
300            // 970 MB
301            Self::ParaphraseMLMpnetBaseV2 | Self::MultilingualE5Base => 970,
302            // 1300 MB
303            Self::BGELargeENV15
304            | Self::BGELargeENV15Q
305            | Self::BGELargeZHV15
306            | Self::MultilingualE5Large
307            | Self::MxbaiEmbedLargeV1
308            | Self::MxbaiEmbedLargeV1Q
309            | Self::GTELargeENV15
310            | Self::GTELargeENV15Q
311            | Self::SnowflakeArcticEmbedL
312            | Self::SnowflakeArcticEmbedLQ => 1300,
313        }
314    }
315
316    /// Original HuggingFace model name (e.g. `"BAAI/bge-small-en-v1.5"`).
317    ///
318    /// This is the **canonical** model identifier from the original publisher.
319    /// For the actual HuggingFace repository used by fastembed-rs for downloads
320    /// and caching, use [`model_code`](Self::model_code) instead.
321    pub const fn model_id(self) -> &'static str {
322        match self {
323            Self::BGESmallENV15 => "BAAI/bge-small-en-v1.5",
324            Self::AllMiniLML6V2 => "sentence-transformers/all-MiniLM-L6-v2",
325            Self::AllMiniLML6V2Q => "Xenova/all-MiniLM-L6-v2",
326            Self::AllMiniLML12V2 => "sentence-transformers/all-MiniLM-L12-v2",
327            Self::AllMiniLML12V2Q => "Xenova/all-MiniLM-L12-v2",
328            Self::AllMpnetBaseV2 => "sentence-transformers/all-mpnet-base-v2",
329            Self::BGEBaseENV15 => "BAAI/bge-base-en-v1.5",
330            Self::BGEBaseENV15Q => "Qdrant/bge-base-en-v1.5-onnx-Q",
331            Self::BGELargeENV15 => "BAAI/bge-large-en-v1.5",
332            Self::BGELargeENV15Q => "Qdrant/bge-large-en-v1.5-onnx-Q",
333            Self::BGESmallENV15Q => "Qdrant/bge-small-en-v1.5-onnx-Q",
334            Self::NomicEmbedTextV1 => "nomic-ai/nomic-embed-text-v1",
335            Self::NomicEmbedTextV15 => "nomic-ai/nomic-embed-text-v1.5",
336            Self::NomicEmbedTextV15Q => "nomic-ai/nomic-embed-text-v1.5",
337            Self::ParaphraseMLMiniLML12V2 => {
338                "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
339            }
340            Self::ParaphraseMLMiniLML12V2Q => "Xenova/paraphrase-multilingual-MiniLM-L12-v2",
341            Self::ParaphraseMLMpnetBaseV2 => {
342                "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
343            }
344            Self::BGESmallZHV15 => "BAAI/bge-small-zh-v1.5",
345            Self::BGELargeZHV15 => "BAAI/bge-large-zh-v1.5",
346            Self::BGEM3 => "BAAI/bge-m3",
347            Self::ModernBertEmbedLarge => "nomic-ai/modernbert-embed-large",
348            Self::MultilingualE5Small => "intfloat/multilingual-e5-small",
349            Self::MultilingualE5Base => "intfloat/multilingual-e5-base",
350            Self::MultilingualE5Large => "intfloat/multilingual-e5-large",
351            Self::MxbaiEmbedLargeV1 => "mixedbread-ai/mxbai-embed-large-v1",
352            Self::MxbaiEmbedLargeV1Q => "mixedbread-ai/mxbai-embed-large-v1",
353            Self::GTEBaseENV15 => "Alibaba-NLP/gte-base-en-v1.5",
354            Self::GTEBaseENV15Q => "Qdrant/gte-base-en-v1.5-onnx-Q",
355            Self::GTELargeENV15 => "Alibaba-NLP/gte-large-en-v1.5",
356            Self::GTELargeENV15Q => "Qdrant/gte-large-en-v1.5-onnx-Q",
357            Self::ClipVitB32 => "openai/clip-vit-base-patch32",
358            Self::JinaEmbeddingsV2BaseCode => "jinaai/jina-embeddings-v2-base-code",
359            Self::JinaEmbeddingsV2BaseEN => "jinaai/jina-embeddings-v2-base-en",
360            Self::EmbeddingGemma300M => "google/embedding-gemma-300M",
361            Self::SnowflakeArcticEmbedXS => "Snowflake/snowflake-arctic-embed-xs",
362            Self::SnowflakeArcticEmbedXSQ => "Snowflake/snowflake-arctic-embed-xs",
363            Self::SnowflakeArcticEmbedS => "Snowflake/snowflake-arctic-embed-s",
364            Self::SnowflakeArcticEmbedSQ => "Snowflake/snowflake-arctic-embed-s",
365            Self::SnowflakeArcticEmbedM => "Snowflake/snowflake-arctic-embed-m",
366            Self::SnowflakeArcticEmbedMQ => "Snowflake/snowflake-arctic-embed-m",
367            Self::SnowflakeArcticEmbedMLong => "Snowflake/snowflake-arctic-embed-m-long",
368            Self::SnowflakeArcticEmbedMLongQ => "Snowflake/snowflake-arctic-embed-m-long",
369            Self::SnowflakeArcticEmbedL => "Snowflake/snowflake-arctic-embed-l",
370            Self::SnowflakeArcticEmbedLQ => "Snowflake/snowflake-arctic-embed-l",
371        }
372    }
373
374    /// Maximum token context per model.
375    pub const fn max_seq_length(self) -> usize {
376        match self {
377            // 256 tokens
378            Self::AllMiniLML6V2
379            | Self::AllMiniLML6V2Q
380            | Self::AllMiniLML12V2
381            | Self::AllMiniLML12V2Q => 256,
382            // 384 tokens
383            Self::AllMpnetBaseV2 => 384,
384            // 8192 tokens
385            Self::BGEM3
386            | Self::NomicEmbedTextV1
387            | Self::NomicEmbedTextV15
388            | Self::NomicEmbedTextV15Q
389            | Self::JinaEmbeddingsV2BaseCode
390            | Self::JinaEmbeddingsV2BaseEN
391            | Self::EmbeddingGemma300M
392            | Self::SnowflakeArcticEmbedMLong
393            | Self::SnowflakeArcticEmbedMLongQ => 8192,
394            // 512 tokens (default)
395            Self::BGESmallENV15
396            | Self::BGESmallENV15Q
397            | Self::BGEBaseENV15
398            | Self::BGEBaseENV15Q
399            | Self::BGELargeENV15
400            | Self::BGELargeENV15Q
401            | Self::BGESmallZHV15
402            | Self::BGELargeZHV15
403            | Self::ParaphraseMLMiniLML12V2
404            | Self::ParaphraseMLMiniLML12V2Q
405            | Self::ParaphraseMLMpnetBaseV2
406            | Self::ModernBertEmbedLarge
407            | Self::MultilingualE5Small
408            | Self::MultilingualE5Base
409            | Self::MultilingualE5Large
410            | Self::MxbaiEmbedLargeV1
411            | Self::MxbaiEmbedLargeV1Q
412            | Self::GTEBaseENV15
413            | Self::GTEBaseENV15Q
414            | Self::GTELargeENV15
415            | Self::GTELargeENV15Q
416            | Self::ClipVitB32
417            | Self::SnowflakeArcticEmbedXS
418            | Self::SnowflakeArcticEmbedXSQ
419            | Self::SnowflakeArcticEmbedS
420            | Self::SnowflakeArcticEmbedSQ
421            | Self::SnowflakeArcticEmbedM
422            | Self::SnowflakeArcticEmbedMQ
423            | Self::SnowflakeArcticEmbedL
424            | Self::SnowflakeArcticEmbedLQ => 512,
425        }
426    }
427
428    /// Whether this is a quantized model (Q suffix).
429    pub const fn is_quantized(self) -> bool {
430        matches!(
431            self,
432            Self::AllMiniLML6V2Q
433                | Self::AllMiniLML12V2Q
434                | Self::BGEBaseENV15Q
435                | Self::BGELargeENV15Q
436                | Self::BGESmallENV15Q
437                | Self::NomicEmbedTextV15Q
438                | Self::ParaphraseMLMiniLML12V2Q
439                | Self::MxbaiEmbedLargeV1Q
440                | Self::GTEBaseENV15Q
441                | Self::GTELargeENV15Q
442                | Self::SnowflakeArcticEmbedXSQ
443                | Self::SnowflakeArcticEmbedSQ
444                | Self::SnowflakeArcticEmbedMQ
445                | Self::SnowflakeArcticEmbedMLongQ
446                | Self::SnowflakeArcticEmbedLQ
447        )
448    }
449
450    /// Whether this model handles image inputs (CLIP).
451    pub const fn is_image_model(self) -> bool {
452        matches!(self, Self::ClipVitB32)
453    }
454
455    /// HuggingFace model code used by fastembed-rs for downloads and caching
456    /// (e.g. `"Xenova/bge-small-en-v1.5"`).
457    ///
458    /// This is the **actual** repository identifier that `hf-hub` uses when
459    /// downloading model weights. Differs from [`model_id`](Self::model_id)
460    /// when the ONNX conversion lives in a separate repo (e.g. `Qdrant/…`,
461    /// `Xenova/…`, `onnx-community/…`).
462    pub const fn model_code(self) -> &'static str {
463        match self {
464            Self::AllMiniLML6V2 => "Qdrant/all-MiniLM-L6-v2-onnx",
465            Self::AllMiniLML6V2Q => "Xenova/all-MiniLM-L6-v2",
466            Self::AllMiniLML12V2 => "Xenova/all-MiniLM-L12-v2",
467            Self::AllMiniLML12V2Q => "Xenova/all-MiniLM-L12-v2",
468            Self::AllMpnetBaseV2 => "Xenova/all-mpnet-base-v2",
469            Self::BGEBaseENV15 => "Xenova/bge-base-en-v1.5",
470            Self::BGEBaseENV15Q => "Qdrant/bge-base-en-v1.5-onnx-Q",
471            Self::BGELargeENV15 => "Xenova/bge-large-en-v1.5",
472            Self::BGELargeENV15Q => "Qdrant/bge-large-en-v1.5-onnx-Q",
473            Self::BGESmallENV15 => "Xenova/bge-small-en-v1.5",
474            Self::BGESmallENV15Q => "Qdrant/bge-small-en-v1.5-onnx-Q",
475            Self::NomicEmbedTextV1 => "nomic-ai/nomic-embed-text-v1",
476            Self::NomicEmbedTextV15 => "nomic-ai/nomic-embed-text-v1.5",
477            Self::NomicEmbedTextV15Q => "nomic-ai/nomic-embed-text-v1.5",
478            Self::ParaphraseMLMiniLML12V2 => "Xenova/paraphrase-multilingual-MiniLM-L12-v2",
479            Self::ParaphraseMLMiniLML12V2Q => "Qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
480            Self::ParaphraseMLMpnetBaseV2 => "Xenova/paraphrase-multilingual-mpnet-base-v2",
481            Self::BGESmallZHV15 => "Xenova/bge-small-zh-v1.5",
482            Self::BGELargeZHV15 => "Xenova/bge-large-zh-v1.5",
483            Self::BGEM3 => "BAAI/bge-m3",
484            Self::ModernBertEmbedLarge => "lightonai/modernbert-embed-large",
485            Self::MultilingualE5Small => "intfloat/multilingual-e5-small",
486            Self::MultilingualE5Base => "intfloat/multilingual-e5-base",
487            Self::MultilingualE5Large => "Qdrant/multilingual-e5-large-onnx",
488            Self::MxbaiEmbedLargeV1 => "mixedbread-ai/mxbai-embed-large-v1",
489            Self::MxbaiEmbedLargeV1Q => "mixedbread-ai/mxbai-embed-large-v1",
490            Self::GTEBaseENV15 => "Alibaba-NLP/gte-base-en-v1.5",
491            Self::GTEBaseENV15Q => "Alibaba-NLP/gte-base-en-v1.5",
492            Self::GTELargeENV15 => "Alibaba-NLP/gte-large-en-v1.5",
493            Self::GTELargeENV15Q => "Alibaba-NLP/gte-large-en-v1.5",
494            Self::ClipVitB32 => "Qdrant/clip-ViT-B-32-text",
495            Self::JinaEmbeddingsV2BaseCode => "jinaai/jina-embeddings-v2-base-code",
496            Self::JinaEmbeddingsV2BaseEN => "jinaai/jina-embeddings-v2-base-en",
497            Self::EmbeddingGemma300M => "onnx-community/embeddinggemma-300m-ONNX",
498            Self::SnowflakeArcticEmbedXS => "snowflake/snowflake-arctic-embed-xs",
499            Self::SnowflakeArcticEmbedXSQ => "snowflake/snowflake-arctic-embed-xs",
500            Self::SnowflakeArcticEmbedS => "snowflake/snowflake-arctic-embed-s",
501            Self::SnowflakeArcticEmbedSQ => "snowflake/snowflake-arctic-embed-s",
502            Self::SnowflakeArcticEmbedM => "Snowflake/snowflake-arctic-embed-m",
503            Self::SnowflakeArcticEmbedMQ => "Snowflake/snowflake-arctic-embed-m",
504            Self::SnowflakeArcticEmbedMLong => "snowflake/snowflake-arctic-embed-m-long",
505            Self::SnowflakeArcticEmbedMLongQ => "snowflake/snowflake-arctic-embed-m-long",
506            Self::SnowflakeArcticEmbedL => "snowflake/snowflake-arctic-embed-l",
507            Self::SnowflakeArcticEmbedLQ => "snowflake/snowflake-arctic-embed-l",
508        }
509    }
510
511    /// String representation matching the enum variant name.
512    pub fn as_str(self) -> &'static str {
513        match self {
514            Self::BGESmallENV15 => "BGESmallENV15",
515            Self::AllMiniLML6V2 => "AllMiniLML6V2",
516            Self::AllMiniLML6V2Q => "AllMiniLML6V2Q",
517            Self::AllMiniLML12V2 => "AllMiniLML12V2",
518            Self::AllMiniLML12V2Q => "AllMiniLML12V2Q",
519            Self::AllMpnetBaseV2 => "AllMpnetBaseV2",
520            Self::BGEBaseENV15 => "BGEBaseENV15",
521            Self::BGEBaseENV15Q => "BGEBaseENV15Q",
522            Self::BGELargeENV15 => "BGELargeENV15",
523            Self::BGELargeENV15Q => "BGELargeENV15Q",
524            Self::BGESmallENV15Q => "BGESmallENV15Q",
525            Self::NomicEmbedTextV1 => "NomicEmbedTextV1",
526            Self::NomicEmbedTextV15 => "NomicEmbedTextV15",
527            Self::NomicEmbedTextV15Q => "NomicEmbedTextV15Q",
528            Self::ParaphraseMLMiniLML12V2 => "ParaphraseMLMiniLML12V2",
529            Self::ParaphraseMLMiniLML12V2Q => "ParaphraseMLMiniLML12V2Q",
530            Self::ParaphraseMLMpnetBaseV2 => "ParaphraseMLMpnetBaseV2",
531            Self::BGESmallZHV15 => "BGESmallZHV15",
532            Self::BGELargeZHV15 => "BGELargeZHV15",
533            Self::BGEM3 => "BGEM3",
534            Self::ModernBertEmbedLarge => "ModernBertEmbedLarge",
535            Self::MultilingualE5Small => "MultilingualE5Small",
536            Self::MultilingualE5Base => "MultilingualE5Base",
537            Self::MultilingualE5Large => "MultilingualE5Large",
538            Self::MxbaiEmbedLargeV1 => "MxbaiEmbedLargeV1",
539            Self::MxbaiEmbedLargeV1Q => "MxbaiEmbedLargeV1Q",
540            Self::GTEBaseENV15 => "GTEBaseENV15",
541            Self::GTEBaseENV15Q => "GTEBaseENV15Q",
542            Self::GTELargeENV15 => "GTELargeENV15",
543            Self::GTELargeENV15Q => "GTELargeENV15Q",
544            Self::ClipVitB32 => "ClipVitB32",
545            Self::JinaEmbeddingsV2BaseCode => "JinaEmbeddingsV2BaseCode",
546            Self::JinaEmbeddingsV2BaseEN => "JinaEmbeddingsV2BaseEN",
547            Self::EmbeddingGemma300M => "EmbeddingGemma300M",
548            Self::SnowflakeArcticEmbedXS => "SnowflakeArcticEmbedXS",
549            Self::SnowflakeArcticEmbedXSQ => "SnowflakeArcticEmbedXSQ",
550            Self::SnowflakeArcticEmbedS => "SnowflakeArcticEmbedS",
551            Self::SnowflakeArcticEmbedSQ => "SnowflakeArcticEmbedSQ",
552            Self::SnowflakeArcticEmbedM => "SnowflakeArcticEmbedM",
553            Self::SnowflakeArcticEmbedMQ => "SnowflakeArcticEmbedMQ",
554            Self::SnowflakeArcticEmbedMLong => "SnowflakeArcticEmbedMLong",
555            Self::SnowflakeArcticEmbedMLongQ => "SnowflakeArcticEmbedMLongQ",
556            Self::SnowflakeArcticEmbedL => "SnowflakeArcticEmbedL",
557            Self::SnowflakeArcticEmbedLQ => "SnowflakeArcticEmbedLQ",
558        }
559    }
560
561    /// Parse a model name (case-insensitive).
562    pub fn parse(s: &str) -> Result<Self, String> {
563        Self::ALL
564            .iter()
565            .find(|m| m.as_str().eq_ignore_ascii_case(s))
566            .copied()
567            .ok_or_else(|| format!("unknown embedding model: {s}"))
568    }
569
570    /// All supported models.
571    pub const ALL: &[Self] = &[
572        Self::BGESmallENV15,
573        Self::AllMiniLML6V2,
574        Self::AllMiniLML6V2Q,
575        Self::AllMiniLML12V2,
576        Self::AllMiniLML12V2Q,
577        Self::AllMpnetBaseV2,
578        Self::BGEBaseENV15,
579        Self::BGEBaseENV15Q,
580        Self::BGELargeENV15,
581        Self::BGELargeENV15Q,
582        Self::BGESmallENV15Q,
583        Self::NomicEmbedTextV1,
584        Self::NomicEmbedTextV15,
585        Self::NomicEmbedTextV15Q,
586        Self::ParaphraseMLMiniLML12V2,
587        Self::ParaphraseMLMiniLML12V2Q,
588        Self::ParaphraseMLMpnetBaseV2,
589        Self::BGESmallZHV15,
590        Self::BGELargeZHV15,
591        Self::BGEM3,
592        Self::ModernBertEmbedLarge,
593        Self::MultilingualE5Small,
594        Self::MultilingualE5Base,
595        Self::MultilingualE5Large,
596        Self::MxbaiEmbedLargeV1,
597        Self::MxbaiEmbedLargeV1Q,
598        Self::GTEBaseENV15,
599        Self::GTEBaseENV15Q,
600        Self::GTELargeENV15,
601        Self::GTELargeENV15Q,
602        Self::ClipVitB32,
603        Self::JinaEmbeddingsV2BaseCode,
604        Self::JinaEmbeddingsV2BaseEN,
605        Self::EmbeddingGemma300M,
606        Self::SnowflakeArcticEmbedXS,
607        Self::SnowflakeArcticEmbedXSQ,
608        Self::SnowflakeArcticEmbedS,
609        Self::SnowflakeArcticEmbedSQ,
610        Self::SnowflakeArcticEmbedM,
611        Self::SnowflakeArcticEmbedMQ,
612        Self::SnowflakeArcticEmbedMLong,
613        Self::SnowflakeArcticEmbedMLongQ,
614        Self::SnowflakeArcticEmbedL,
615        Self::SnowflakeArcticEmbedLQ,
616    ];
617
618    /// Map to `fastembed::EmbeddingModel`.
619    ///
620    /// Only available when an ONNX embedding feature is enabled
621    /// (`embedding-fastembed` or `embedding-fastembed-dynamic-linking`).
622    #[cfg(any(
623        feature = "embedding-fastembed",
624        feature = "embedding-fastembed-dynamic-linking"
625    ))]
626    pub fn as_fastembed(self) -> fastembed::EmbeddingModel {
627        match self {
628            Self::BGESmallENV15 => fastembed::EmbeddingModel::BGESmallENV15,
629            Self::AllMiniLML6V2 => fastembed::EmbeddingModel::AllMiniLML6V2,
630            Self::AllMiniLML6V2Q => fastembed::EmbeddingModel::AllMiniLML6V2Q,
631            Self::AllMiniLML12V2 => fastembed::EmbeddingModel::AllMiniLML12V2,
632            Self::AllMiniLML12V2Q => fastembed::EmbeddingModel::AllMiniLML12V2Q,
633            Self::AllMpnetBaseV2 => fastembed::EmbeddingModel::AllMpnetBaseV2,
634            Self::BGEBaseENV15 => fastembed::EmbeddingModel::BGEBaseENV15,
635            Self::BGEBaseENV15Q => fastembed::EmbeddingModel::BGEBaseENV15Q,
636            Self::BGELargeENV15 => fastembed::EmbeddingModel::BGELargeENV15,
637            Self::BGELargeENV15Q => fastembed::EmbeddingModel::BGELargeENV15Q,
638            Self::BGESmallENV15Q => fastembed::EmbeddingModel::BGESmallENV15Q,
639            Self::NomicEmbedTextV1 => fastembed::EmbeddingModel::NomicEmbedTextV1,
640            Self::NomicEmbedTextV15 => fastembed::EmbeddingModel::NomicEmbedTextV15,
641            Self::NomicEmbedTextV15Q => fastembed::EmbeddingModel::NomicEmbedTextV15Q,
642            Self::ParaphraseMLMiniLML12V2 => fastembed::EmbeddingModel::ParaphraseMLMiniLML12V2,
643            Self::ParaphraseMLMiniLML12V2Q => fastembed::EmbeddingModel::ParaphraseMLMiniLML12V2Q,
644            Self::ParaphraseMLMpnetBaseV2 => fastembed::EmbeddingModel::ParaphraseMLMpnetBaseV2,
645            Self::BGESmallZHV15 => fastembed::EmbeddingModel::BGESmallZHV15,
646            Self::BGELargeZHV15 => fastembed::EmbeddingModel::BGELargeZHV15,
647            Self::BGEM3 => fastembed::EmbeddingModel::BGEM3,
648            Self::ModernBertEmbedLarge => fastembed::EmbeddingModel::ModernBertEmbedLarge,
649            Self::MultilingualE5Small => fastembed::EmbeddingModel::MultilingualE5Small,
650            Self::MultilingualE5Base => fastembed::EmbeddingModel::MultilingualE5Base,
651            Self::MultilingualE5Large => fastembed::EmbeddingModel::MultilingualE5Large,
652            Self::MxbaiEmbedLargeV1 => fastembed::EmbeddingModel::MxbaiEmbedLargeV1,
653            Self::MxbaiEmbedLargeV1Q => fastembed::EmbeddingModel::MxbaiEmbedLargeV1Q,
654            Self::GTEBaseENV15 => fastembed::EmbeddingModel::GTEBaseENV15,
655            Self::GTEBaseENV15Q => fastembed::EmbeddingModel::GTEBaseENV15Q,
656            Self::GTELargeENV15 => fastembed::EmbeddingModel::GTELargeENV15,
657            Self::GTELargeENV15Q => fastembed::EmbeddingModel::GTELargeENV15Q,
658            Self::ClipVitB32 => fastembed::EmbeddingModel::ClipVitB32,
659            Self::JinaEmbeddingsV2BaseCode => fastembed::EmbeddingModel::JinaEmbeddingsV2BaseCode,
660            Self::JinaEmbeddingsV2BaseEN => fastembed::EmbeddingModel::JinaEmbeddingsV2BaseEN,
661            Self::EmbeddingGemma300M => fastembed::EmbeddingModel::EmbeddingGemma300M,
662            Self::SnowflakeArcticEmbedXS => fastembed::EmbeddingModel::SnowflakeArcticEmbedXS,
663            Self::SnowflakeArcticEmbedXSQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedXSQ,
664            Self::SnowflakeArcticEmbedS => fastembed::EmbeddingModel::SnowflakeArcticEmbedS,
665            Self::SnowflakeArcticEmbedSQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedSQ,
666            Self::SnowflakeArcticEmbedM => fastembed::EmbeddingModel::SnowflakeArcticEmbedM,
667            Self::SnowflakeArcticEmbedMQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedMQ,
668            Self::SnowflakeArcticEmbedMLong => fastembed::EmbeddingModel::SnowflakeArcticEmbedMLong,
669            Self::SnowflakeArcticEmbedMLongQ => {
670                fastembed::EmbeddingModel::SnowflakeArcticEmbedMLongQ
671            }
672            Self::SnowflakeArcticEmbedL => fastembed::EmbeddingModel::SnowflakeArcticEmbedL,
673            Self::SnowflakeArcticEmbedLQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedLQ,
674        }
675    }
676}
677
678impl std::fmt::Display for EmbeddingModel {
679    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
680        f.write_str(self.as_str())
681    }
682}
683
684impl std::str::FromStr for EmbeddingModel {
685    type Err = String;
686
687    fn from_str(s: &str) -> Result<Self, Self::Err> {
688        Self::parse(s)
689    }
690}
691
692#[cfg(test)]
693mod tests {
694    use super::*;
695
696    #[test]
697    fn all_count() {
698        assert_eq!(EmbeddingModel::ALL.len(), 44);
699    }
700
701    #[test]
702    fn default_is_bge_small() {
703        assert_eq!(EmbeddingModel::default(), EmbeddingModel::BGESmallENV15);
704    }
705
706    #[test]
707    fn dimension_consistency() {
708        for &m in EmbeddingModel::ALL {
709            let dim = m.dimension();
710            assert!(
711                [384, 512, 768, 1024].contains(&dim),
712                "{m:?}: unexpected dimension {dim}"
713            );
714        }
715    }
716
717    #[test]
718    fn parse_roundtrip() {
719        for &m in EmbeddingModel::ALL {
720            let s = m.as_str();
721            assert_eq!(EmbeddingModel::parse(s).unwrap(), m);
722        }
723    }
724
725    #[test]
726    fn parse_case_insensitive() {
727        assert_eq!(
728            EmbeddingModel::parse("bgesmallenv15").unwrap(),
729            EmbeddingModel::BGESmallENV15
730        );
731        assert_eq!(
732            EmbeddingModel::parse("ALLMINILML6V2").unwrap(),
733            EmbeddingModel::AllMiniLML6V2
734        );
735    }
736
737    #[test]
738    fn parse_unknown_fails() {
739        assert!(EmbeddingModel::parse("NotARealModel").is_err());
740    }
741
742    #[test]
743    fn quantized_flags() {
744        let quantized: Vec<_> = EmbeddingModel::ALL
745            .iter()
746            .filter(|m| m.is_quantized())
747            .copied()
748            .collect();
749        // All Q-suffix variants should be flagged
750        for m in &quantized {
751            assert!(m.as_str().ends_with('Q'), "{m:?} flagged but no Q suffix");
752        }
753        // Non-Q variants should NOT be flagged
754        for &m in EmbeddingModel::ALL {
755            if !m.as_str().ends_with('Q') {
756                assert!(!m.is_quantized(), "{m:?} not Q but flagged quantized");
757            }
758        }
759    }
760
761    #[test]
762    fn image_model_flag() {
763        assert!(EmbeddingModel::ClipVitB32.is_image_model());
764        assert_eq!(
765            EmbeddingModel::ALL
766                .iter()
767                .filter(|m| m.is_image_model())
768                .count(),
769            1
770        );
771    }
772
773    #[test]
774    fn prefix_mapping() {
775        // E5 models have query + doc prefixes
776        for &m in &[
777            EmbeddingModel::MultilingualE5Small,
778            EmbeddingModel::MultilingualE5Base,
779            EmbeddingModel::MultilingualE5Large,
780        ] {
781            assert_eq!(m.query_prefix(), Some("query: "));
782            assert_eq!(m.doc_prefix(), Some("passage: "));
783        }
784        // Snowflake models have query prefix only
785        for &m in &[
786            EmbeddingModel::SnowflakeArcticEmbedXS,
787            EmbeddingModel::SnowflakeArcticEmbedLQ,
788        ] {
789            assert!(m.query_prefix().is_some());
790            assert!(m.doc_prefix().is_none());
791        }
792        // Nomic v1.5 models have query + doc prefixes
793        for &m in &[
794            EmbeddingModel::NomicEmbedTextV15,
795            EmbeddingModel::NomicEmbedTextV15Q,
796        ] {
797            assert_eq!(m.query_prefix(), Some("search_query: "));
798            assert_eq!(m.doc_prefix(), Some("search_document: "));
799        }
800        // Nomic v1 has no prefixes
801        assert!(EmbeddingModel::NomicEmbedTextV1.query_prefix().is_none());
802        assert!(EmbeddingModel::NomicEmbedTextV1.doc_prefix().is_none());
803        // Most models have no prefixes
804        assert!(EmbeddingModel::BGESmallENV15.query_prefix().is_none());
805        assert!(EmbeddingModel::BGESmallENV15.doc_prefix().is_none());
806    }
807
808    #[test]
809    fn from_str_trait() {
810        let m: EmbeddingModel = "BGESmallENV15".parse().unwrap();
811        assert_eq!(m, EmbeddingModel::BGESmallENV15);
812    }
813
814    #[test]
815    fn display_trait() {
816        assert_eq!(EmbeddingModel::BGESmallENV15.to_string(), "BGESmallENV15");
817    }
818
819    #[test]
820    fn metadata_nonzero() {
821        for &m in EmbeddingModel::ALL {
822            assert!(m.size_mb() > 0, "{m:?}: size_mb is zero");
823            assert!(!m.model_id().is_empty(), "{m:?}: model_id is empty");
824            assert!(!m.model_code().is_empty(), "{m:?}: model_code is empty");
825            assert!(m.max_seq_length() > 0, "{m:?}: max_seq_length is zero");
826        }
827    }
828
829    #[test]
830    fn max_seq_length_values() {
831        assert_eq!(EmbeddingModel::AllMiniLML6V2.max_seq_length(), 256);
832        assert_eq!(EmbeddingModel::AllMpnetBaseV2.max_seq_length(), 384);
833        assert_eq!(EmbeddingModel::BGEM3.max_seq_length(), 8192);
834        assert_eq!(EmbeddingModel::BGESmallENV15.max_seq_length(), 512);
835    }
836}
llm_kernel/embedding/catalog.rs

llm_kernel/embedding/
catalog.rs
