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//! # Backend availability
7//!
8//! Each model's backend support is independent — a model listed below is an
9//! ONNX checkpoint served via `embedding-fastembed`; the same model family
10//! may also have a Rust-native path on specific platforms:
11//!
12//! | Backend (feature)            | Models                                       | Platform         |
13//! |------------------------------|----------------------------------------------|------------------|
14//! | `embedding-fastembed` (ONNX) | All 44 variants below                        | cross-platform   |
15//! | `embedding-metal` (candle)   | Qwen3-Embedding, Nomic V2 MoE                | macOS (Metal)    |
16//! | `embedding-mlx` (MLX)        | 13 vanilla-BERT models (see `mlx_supported`)  | macOS (aarch64)  |
17//!
18//! The MLX provider (`MlxEmbeddingProvider`, feature `embedding-mlx`) covers the
19//! vanilla-BERT encoder family: BGE-en-v1.5 (small/base/large), bge-small-zh-v1.5,
20//! all-MiniLM-L6/L12, paraphrase-multilingual-MiniLM, multilingual-e5-small,
21//! Snowflake Arctic (xs/s/m/l) and mxbai-embed-large. Other architectures
22//! (XLM-R, MPNet, NomicBert, GTE `NewModel`, JinaBert, ModernBERT, Gemma, CLIP)
23//! need their own forward passes. See [`EmbeddingModel::mlx_supported`] for the
24//! exclusion list and `src/embedding/mlx.rs`.
25
26/// Embedding model catalog with metadata for all supported ONNX models.
27///
28/// Variant names match `fastembed::EmbeddingModel` exactly for trivial 1:1 mapping.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
30pub enum EmbeddingModel {
31    // ── sentence-transformers ───────────────────────
32    /// BGE Small EN v1.5 — fast 384-dim English model (default).
33    ///
34    /// One of the vanilla-BERT models served by the Rust-native MLX backend
35    /// (`embedding-mlx`, macOS aarch64) — see [`Self::mlx_supported`].
36    #[default]
37    BGESmallENV15,
38    /// sentence-transformers all-MiniLM-L6-v2 (384-dim).
39    AllMiniLML6V2,
40    /// Quantized all-MiniLM-L6-v2 (384-dim).
41    AllMiniLML6V2Q,
42    /// sentence-transformers all-MiniLM-L12-v2 (384-dim).
43    AllMiniLML12V2,
44    /// Quantized all-MiniLM-L12-v2 (384-dim).
45    AllMiniLML12V2Q,
46    /// sentence-transformers all-mpnet-base-v2 (768-dim).
47    AllMpnetBaseV2,
48    // ── BGE family ──────────────────────────────────
49    /// BGE Base EN v1.5 (768-dim).
50    BGEBaseENV15,
51    /// Quantized BGE Base EN v1.5 (768-dim).
52    BGEBaseENV15Q,
53    /// BGE Large EN v1.5 (1024-dim).
54    BGELargeENV15,
55    /// Quantized BGE Large EN v1.5 (1024-dim).
56    BGELargeENV15Q,
57    /// Quantized BGE Small EN v1.5 (384-dim).
58    BGESmallENV15Q,
59    /// BGE Small ZH v1.5 — Chinese model (512-dim).
60    BGESmallZHV15,
61    /// BGE Large ZH v1.5 — Chinese model (1024-dim).
62    BGELargeZHV15,
63    /// BGE M3 — multilingual, 100+ languages, 8192 context (1024-dim).
64    BGEM3,
65    // ── Nomic ───────────────────────────────────────
66    /// Nomic Embed Text v1 — 8192 context English model (768-dim).
67    NomicEmbedTextV1,
68    /// Nomic Embed Text v1.5 — 8192 context English model (768-dim).
69    NomicEmbedTextV15,
70    /// Quantized Nomic Embed Text v1.5 (768-dim).
71    NomicEmbedTextV15Q,
72    // ── Paraphrase ──────────────────────────────────
73    /// Paraphrase multilingual MiniLM-L12-v2 (384-dim).
74    ParaphraseMLMiniLML12V2,
75    /// Quantized paraphrase multilingual MiniLM-L12-v2 (384-dim).
76    ParaphraseMLMiniLML12V2Q,
77    /// Paraphrase multilingual mpnet-base-v2 (768-dim).
78    ParaphraseMLMpnetBaseV2,
79    // ── ModernBERT ──────────────────────────────────
80    /// ModernBERT Embed Large (1024-dim).
81    ModernBertEmbedLarge,
82    // ── E5 multilingual ─────────────────────────────
83    /// Multilingual E5 Small (384-dim).
84    MultilingualE5Small,
85    /// Multilingual E5 Base (768-dim).
86    MultilingualE5Base,
87    /// Multilingual E5 Large (1024-dim).
88    MultilingualE5Large,
89    // ── Mixedbread ──────────────────────────────────
90    /// Mixedbread mxbai-embed-large-v1 (1024-dim).
91    MxbaiEmbedLargeV1,
92    /// Quantized mxbai-embed-large-v1 (1024-dim).
93    MxbaiEmbedLargeV1Q,
94    // ── GTE (Alibaba) ──────────────────────────────
95    /// Alibaba GTE Base EN v1.5 (768-dim).
96    GTEBaseENV15,
97    /// Quantized Alibaba GTE Base EN v1.5 (768-dim).
98    GTEBaseENV15Q,
99    /// Alibaba GTE Large EN v1.5 (1024-dim).
100    GTELargeENV15,
101    /// Quantized Alibaba GTE Large EN v1.5 (1024-dim).
102    GTELargeENV15Q,
103    // ── CLIP ────────────────────────────────────────
104    /// OpenAI CLIP ViT-B/32 text encoder (512-dim).
105    ClipVitB32,
106    // ── Jina ────────────────────────────────────────
107    /// Jina Embeddings v2 Base for code (768-dim).
108    JinaEmbeddingsV2BaseCode,
109    /// Jina Embeddings v2 Base for English (768-dim).
110    JinaEmbeddingsV2BaseEN,
111    // ── Gemma ───────────────────────────────────────
112    /// Google EmbeddingGemma 300M (768-dim).
113    EmbeddingGemma300M,
114    // ── Snowflake Arctic ────────────────────────────
115    /// Snowflake Arctic Embed XS (384-dim).
116    SnowflakeArcticEmbedXS,
117    /// Quantized Snowflake Arctic Embed XS (384-dim).
118    SnowflakeArcticEmbedXSQ,
119    /// Snowflake Arctic Embed S (384-dim).
120    SnowflakeArcticEmbedS,
121    /// Quantized Snowflake Arctic Embed S (384-dim).
122    SnowflakeArcticEmbedSQ,
123    /// Snowflake Arctic Embed M (768-dim).
124    SnowflakeArcticEmbedM,
125    /// Quantized Snowflake Arctic Embed M (768-dim).
126    SnowflakeArcticEmbedMQ,
127    /// Snowflake Arctic Embed M Long — 2048 context (768-dim).
128    SnowflakeArcticEmbedMLong,
129    /// Quantized Snowflake Arctic Embed M Long (768-dim).
130    SnowflakeArcticEmbedMLongQ,
131    /// Snowflake Arctic Embed L (1024-dim).
132    SnowflakeArcticEmbedL,
133    /// Quantized Snowflake Arctic Embed L (1024-dim).
134    SnowflakeArcticEmbedLQ,
135}
136
137impl EmbeddingModel {
138    /// Embedding dimensionality.
139    pub const fn dimension(self) -> usize {
140        match self {
141            // 384-dim
142            Self::BGESmallENV15
143            | Self::AllMiniLML6V2
144            | Self::AllMiniLML6V2Q
145            | Self::AllMiniLML12V2
146            | Self::AllMiniLML12V2Q
147            | Self::BGESmallENV15Q
148            | Self::ParaphraseMLMiniLML12V2
149            | Self::ParaphraseMLMiniLML12V2Q
150            | Self::MultilingualE5Small
151            | Self::SnowflakeArcticEmbedXS
152            | Self::SnowflakeArcticEmbedXSQ
153            | Self::SnowflakeArcticEmbedS
154            | Self::SnowflakeArcticEmbedSQ => 384,
155            // 512-dim
156            Self::BGESmallZHV15 | Self::ClipVitB32 => 512,
157            // 768-dim
158            Self::AllMpnetBaseV2
159            | Self::BGEBaseENV15
160            | Self::BGEBaseENV15Q
161            | Self::NomicEmbedTextV1
162            | Self::NomicEmbedTextV15
163            | Self::NomicEmbedTextV15Q
164            | Self::ParaphraseMLMpnetBaseV2
165            | Self::MultilingualE5Base
166            | Self::GTEBaseENV15
167            | Self::GTEBaseENV15Q
168            | Self::JinaEmbeddingsV2BaseCode
169            | Self::JinaEmbeddingsV2BaseEN
170            | Self::EmbeddingGemma300M
171            | Self::SnowflakeArcticEmbedM
172            | Self::SnowflakeArcticEmbedMQ
173            | Self::SnowflakeArcticEmbedMLong
174            | Self::SnowflakeArcticEmbedMLongQ => 768,
175            // 1024-dim
176            Self::BGELargeENV15
177            | Self::BGELargeENV15Q
178            | Self::BGELargeZHV15
179            | Self::BGEM3
180            | Self::ModernBertEmbedLarge
181            | Self::MultilingualE5Large
182            | Self::MxbaiEmbedLargeV1
183            | Self::MxbaiEmbedLargeV1Q
184            | Self::GTELargeENV15
185            | Self::GTELargeENV15Q
186            | Self::SnowflakeArcticEmbedL
187            | Self::SnowflakeArcticEmbedLQ => 1024,
188        }
189    }
190
191    /// Short human-readable description.
192    pub const fn description(self) -> &'static str {
193        match self {
194            Self::BGESmallENV15 => "v1.5 release of the fast and default English model",
195            Self::AllMiniLML6V2 => "Sentence Transformer model, MiniLM-L6-v2",
196            Self::AllMiniLML6V2Q => "Quantized Sentence Transformer model, MiniLM-L6-v2",
197            Self::AllMiniLML12V2 => "Sentence Transformer model, MiniLM-L12-v2",
198            Self::AllMiniLML12V2Q => "Quantized Sentence Transformer model, MiniLM-L12-v2",
199            Self::AllMpnetBaseV2 => "Sentence Transformer model, mpnet-base-v2",
200            Self::BGEBaseENV15 => "v1.5 release of the base English model",
201            Self::BGEBaseENV15Q => "Quantized v1.5 release of the base English model",
202            Self::BGELargeENV15 => "v1.5 release of the large English model",
203            Self::BGELargeENV15Q => "Quantized v1.5 release of the large English model",
204            Self::BGESmallENV15Q => "Quantized v1.5 release of the fast and default English model",
205            Self::NomicEmbedTextV1 => "8192 context length english model",
206            Self::NomicEmbedTextV15 => "v1.5 release of the 8192 context length english model",
207            Self::NomicEmbedTextV15Q => {
208                "Quantized v1.5 release of the 8192 context length english model"
209            }
210            Self::ParaphraseMLMiniLML12V2 => "Multi-lingual model",
211            Self::ParaphraseMLMiniLML12V2Q => "Quantized multi-lingual model",
212            Self::ParaphraseMLMpnetBaseV2 => {
213                "Sentence-transformers model for clustering or semantic search"
214            }
215            Self::BGESmallZHV15 => "v1.5 release of the small Chinese model",
216            Self::BGELargeZHV15 => "v1.5 release of the large Chinese model",
217            Self::BGEM3 => "Multilingual M3 model with 8192 context length, 100+ languages",
218            Self::ModernBertEmbedLarge => "Large model of ModernBert Text Embeddings",
219            Self::MultilingualE5Small => "Small model of multilingual E5 Text Embeddings",
220            Self::MultilingualE5Base => "Base model of multilingual E5 Text Embeddings",
221            Self::MultilingualE5Large => "Large model of multilingual E5 Text Embeddings",
222            Self::MxbaiEmbedLargeV1 => "Large English embedding model from MixedBreed.ai",
223            Self::MxbaiEmbedLargeV1Q => {
224                "Quantized large English embedding model from MixedBreed.ai"
225            }
226            Self::GTEBaseENV15 => "Base multilingual embedding model from Alibaba",
227            Self::GTEBaseENV15Q => "Quantized base multilingual embedding model from Alibaba",
228            Self::GTELargeENV15 => "Large multilingual embedding model from Alibaba",
229            Self::GTELargeENV15Q => "Quantized large multilingual embedding model from Alibaba",
230            Self::ClipVitB32 => "CLIP text encoder based on ViT-B/32",
231            Self::JinaEmbeddingsV2BaseCode => "Jina embeddings v2 base code",
232            Self::JinaEmbeddingsV2BaseEN => "Jina embeddings v2 base English",
233            Self::EmbeddingGemma300M => "EmbeddingGemma 300M parameter model from Google",
234            Self::SnowflakeArcticEmbedXS => "Snowflake Arctic embed model, xs",
235            Self::SnowflakeArcticEmbedXSQ => "Quantized Snowflake Arctic embed model, xs",
236            Self::SnowflakeArcticEmbedS => "Snowflake Arctic embed model, small",
237            Self::SnowflakeArcticEmbedSQ => "Quantized Snowflake Arctic embed model, small",
238            Self::SnowflakeArcticEmbedM => "Snowflake Arctic embed model, medium",
239            Self::SnowflakeArcticEmbedMQ => "Quantized Snowflake Arctic embed model, medium",
240            Self::SnowflakeArcticEmbedMLong => {
241                "Snowflake Arctic embed model, medium with 2048 context"
242            }
243            Self::SnowflakeArcticEmbedMLongQ => {
244                "Quantized Snowflake Arctic embed model, medium with 2048 context"
245            }
246            Self::SnowflakeArcticEmbedL => "Snowflake Arctic embed model, large",
247            Self::SnowflakeArcticEmbedLQ => "Quantized Snowflake Arctic embed model, large",
248        }
249    }
250
251    /// Optional prefix prepended to query texts before embedding.
252    pub const fn query_prefix(self) -> Option<&'static str> {
253        match self {
254            Self::MultilingualE5Small | Self::MultilingualE5Base | Self::MultilingualE5Large => {
255                Some("query: ")
256            }
257            Self::NomicEmbedTextV15 | Self::NomicEmbedTextV15Q => Some("search_query: "),
258            // Asymmetric models using the BGE instruction prefix.
259            Self::BGESmallENV15
260            | Self::BGESmallENV15Q
261            | Self::BGEBaseENV15
262            | Self::BGEBaseENV15Q
263            | Self::BGELargeENV15
264            | Self::BGELargeENV15Q
265            | Self::MxbaiEmbedLargeV1
266            | Self::MxbaiEmbedLargeV1Q
267            | Self::SnowflakeArcticEmbedXS
268            | Self::SnowflakeArcticEmbedXSQ
269            | Self::SnowflakeArcticEmbedS
270            | Self::SnowflakeArcticEmbedSQ
271            | Self::SnowflakeArcticEmbedM
272            | Self::SnowflakeArcticEmbedMQ
273            | Self::SnowflakeArcticEmbedMLong
274            | Self::SnowflakeArcticEmbedMLongQ
275            | Self::SnowflakeArcticEmbedL
276            | Self::SnowflakeArcticEmbedLQ => {
277                Some("Represent this sentence for searching relevant passages: ")
278            }
279            _ => None,
280        }
281    }
282
283    /// Optional prefix prepended to document texts before embedding.
284    pub const fn doc_prefix(self) -> Option<&'static str> {
285        match self {
286            Self::MultilingualE5Small | Self::MultilingualE5Base | Self::MultilingualE5Large => {
287                Some("passage: ")
288            }
289            Self::NomicEmbedTextV15 | Self::NomicEmbedTextV15Q => Some("search_document: "),
290            _ => None,
291        }
292    }
293
294    /// Whether the MLX Rust-native backend (`embedding-mlx`) supports this model.
295    ///
296    /// True only for models whose **original weight repo** is a vanilla-BERT
297    /// encoder the MLX forward pass can actually load: `architectures:
298    /// ["BertModel"]`, `model_type: "bert"`, absolute position embeddings, gelu,
299    /// a `model.safetensors` in F32/F16/BF16, and the standard
300    /// `encoder.layer.N.attention.self.*` tensor layout. Each entry below was
301    /// verified by probing that repo's `config.json` and safetensors header.
302    ///
303    /// Excluded despite being in the catalog:
304    /// - `SnowflakeArcticEmbedMLong`, `NomicEmbedTextV1/V15` — `NomicBertModel`
305    ///   (`encoder.layers.N.attn.Wqkv`, no position embeddings)
306    /// - `GTEBaseENV15`, `GTELargeENV15` — `NewModel`
307    /// - `MultilingualE5Base/Large`, `ParaphraseMLMpnetBaseV2`, `BGEM3` —
308    ///   `XLMRobertaModel`
309    /// - `AllMpnetBaseV2` — `MPNetForMaskedLM`
310    /// - `JinaEmbeddingsV2Base*` — `JinaBertForMaskedLM` (ALiBi)
311    /// - `ModernBertEmbedLarge` — `ModernBertModel`
312    /// - `EmbeddingGemma300M`, `ClipVitB32` — not BERT text encoders
313    /// - `BGELargeZHV15` — BertModel, but the repo ships only
314    ///   `pytorch_model.bin` with no safetensors to load
315    pub const fn mlx_supported(self) -> bool {
316        matches!(
317            self,
318            Self::BGESmallENV15
319                | Self::BGESmallENV15Q
320                | Self::BGEBaseENV15
321                | Self::BGEBaseENV15Q
322                | Self::BGELargeENV15
323                | Self::BGELargeENV15Q
324                | Self::BGESmallZHV15
325                | Self::AllMiniLML6V2
326                | Self::AllMiniLML6V2Q
327                | Self::AllMiniLML12V2
328                | Self::AllMiniLML12V2Q
329                | Self::ParaphraseMLMiniLML12V2
330                | Self::ParaphraseMLMiniLML12V2Q
331                | Self::MultilingualE5Small
332                | Self::SnowflakeArcticEmbedXS
333                | Self::SnowflakeArcticEmbedXSQ
334                | Self::SnowflakeArcticEmbedS
335                | Self::SnowflakeArcticEmbedSQ
336                | Self::SnowflakeArcticEmbedM
337                | Self::SnowflakeArcticEmbedMQ
338                | Self::SnowflakeArcticEmbedL
339                | Self::SnowflakeArcticEmbedLQ
340                | Self::MxbaiEmbedLargeV1
341                | Self::MxbaiEmbedLargeV1Q
342        )
343    }
344
345    /// Pooling strategy: `true` = CLS (token 0), `false` = mean (mask-weighted).
346    ///
347    /// Read from each model's `1_Pooling/config.json`: BGE / Arctic / mxbai set
348    /// `pooling_mode_cls_token`; MiniLM, paraphrase-multilingual-MiniLM and
349    /// multilingual-e5-small set `pooling_mode_mean_tokens`.
350    pub const fn uses_cls_pooling(self) -> bool {
351        matches!(
352            self,
353            Self::BGESmallENV15
354                | Self::BGESmallENV15Q
355                | Self::BGEBaseENV15
356                | Self::BGEBaseENV15Q
357                | Self::BGELargeENV15
358                | Self::BGELargeENV15Q
359                | Self::BGESmallZHV15
360                | Self::BGELargeZHV15
361                | Self::SnowflakeArcticEmbedXS
362                | Self::SnowflakeArcticEmbedXSQ
363                | Self::SnowflakeArcticEmbedS
364                | Self::SnowflakeArcticEmbedSQ
365                | Self::SnowflakeArcticEmbedM
366                | Self::SnowflakeArcticEmbedMQ
367                | Self::SnowflakeArcticEmbedMLong
368                | Self::SnowflakeArcticEmbedMLongQ
369                | Self::SnowflakeArcticEmbedL
370                | Self::SnowflakeArcticEmbedLQ
371                | Self::MxbaiEmbedLargeV1
372                | Self::MxbaiEmbedLargeV1Q
373        )
374    }
375
376    /// HuggingFace repo holding the **original PyTorch/safetensors weights** for
377    /// the MLX backend. Differs from [`model_code`](Self::model_code) (which
378    /// points at the ONNX-converted repo) for BGE/MiniLM/paraphrase models —
379    /// MLX loads `model.safetensors` + `config.json`, which the `Xenova/*` /
380    /// `Qdrant/*` ONNX repos do not carry.
381    ///
382    /// Only meaningful when [`mlx_supported`](Self::mlx_supported) is true.
383    pub const fn mlx_repo(self) -> &'static str {
384        match self {
385            Self::BGESmallENV15 | Self::BGESmallENV15Q => "BAAI/bge-small-en-v1.5",
386            Self::BGEBaseENV15 | Self::BGEBaseENV15Q => "BAAI/bge-base-en-v1.5",
387            Self::BGELargeENV15 | Self::BGELargeENV15Q => "BAAI/bge-large-en-v1.5",
388            Self::BGESmallZHV15 => "BAAI/bge-small-zh-v1.5",
389            Self::AllMiniLML6V2 | Self::AllMiniLML6V2Q => "sentence-transformers/all-MiniLM-L6-v2",
390            Self::AllMiniLML12V2 | Self::AllMiniLML12V2Q => {
391                "sentence-transformers/all-MiniLM-L12-v2"
392            }
393            Self::ParaphraseMLMiniLML12V2 | Self::ParaphraseMLMiniLML12V2Q => {
394                "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
395            }
396            // Arctic / mxbai / e5-small: model_code already points at the
397            // original (non-ONNX) repo, which carries model.safetensors.
398            _ => self.model_code(),
399        }
400    }
401
402    /// Approximate ONNX model size in MB.
403    pub const fn size_mb(self) -> usize {
404        match self {
405            // 40 MB
406            Self::BGESmallENV15 | Self::BGESmallENV15Q => 40,
407            // 80 MB
408            Self::AllMiniLML6V2 | Self::AllMiniLML6V2Q => 80,
409            // 90 MB
410            Self::SnowflakeArcticEmbedXS | Self::SnowflakeArcticEmbedXSQ | Self::BGESmallZHV15 => {
411                90
412            }
413            // 120 MB
414            Self::AllMiniLML12V2 | Self::AllMiniLML12V2Q => 120,
415            // 130 MB
416            Self::SnowflakeArcticEmbedS | Self::SnowflakeArcticEmbedSQ => 130,
417            // 260 MB
418            Self::JinaEmbeddingsV2BaseCode
419            | Self::JinaEmbeddingsV2BaseEN
420            | Self::EmbeddingGemma300M => 260,
421            // 420 MB
422            Self::AllMpnetBaseV2
423            | Self::BGEBaseENV15
424            | Self::BGEBaseENV15Q
425            | Self::GTEBaseENV15
426            | Self::GTEBaseENV15Q => 420,
427            // 430 MB
428            Self::SnowflakeArcticEmbedM
429            | Self::SnowflakeArcticEmbedMQ
430            | Self::SnowflakeArcticEmbedMLong
431            | Self::SnowflakeArcticEmbedMLongQ => 430,
432            // 470 MB
433            Self::ParaphraseMLMiniLML12V2
434            | Self::ParaphraseMLMiniLML12V2Q
435            | Self::MultilingualE5Small => 470,
436            // 550 MB
437            Self::NomicEmbedTextV1 | Self::NomicEmbedTextV15 | Self::NomicEmbedTextV15Q => 550,
438            // 600 MB
439            Self::BGEM3 | Self::ModernBertEmbedLarge | Self::ClipVitB32 => 600,
440            // 970 MB
441            Self::ParaphraseMLMpnetBaseV2 | Self::MultilingualE5Base => 970,
442            // 1300 MB
443            Self::BGELargeENV15
444            | Self::BGELargeENV15Q
445            | Self::BGELargeZHV15
446            | Self::MultilingualE5Large
447            | Self::MxbaiEmbedLargeV1
448            | Self::MxbaiEmbedLargeV1Q
449            | Self::GTELargeENV15
450            | Self::GTELargeENV15Q
451            | Self::SnowflakeArcticEmbedL
452            | Self::SnowflakeArcticEmbedLQ => 1300,
453        }
454    }
455
456    /// Original HuggingFace model name (e.g. `"BAAI/bge-small-en-v1.5"`).
457    ///
458    /// This is the **canonical** model identifier from the original publisher.
459    /// For the actual HuggingFace repository used by fastembed-rs for downloads
460    /// and caching, use [`model_code`](Self::model_code) instead.
461    pub const fn model_id(self) -> &'static str {
462        match self {
463            Self::BGESmallENV15 => "BAAI/bge-small-en-v1.5",
464            Self::AllMiniLML6V2 => "sentence-transformers/all-MiniLM-L6-v2",
465            Self::AllMiniLML6V2Q => "Xenova/all-MiniLM-L6-v2",
466            Self::AllMiniLML12V2 => "sentence-transformers/all-MiniLM-L12-v2",
467            Self::AllMiniLML12V2Q => "Xenova/all-MiniLM-L12-v2",
468            Self::AllMpnetBaseV2 => "sentence-transformers/all-mpnet-base-v2",
469            Self::BGEBaseENV15 => "BAAI/bge-base-en-v1.5",
470            Self::BGEBaseENV15Q => "Qdrant/bge-base-en-v1.5-onnx-Q",
471            Self::BGELargeENV15 => "BAAI/bge-large-en-v1.5",
472            Self::BGELargeENV15Q => "Qdrant/bge-large-en-v1.5-onnx-Q",
473            Self::BGESmallENV15Q => "Qdrant/bge-small-en-v1.5-onnx-Q",
474            Self::NomicEmbedTextV1 => "nomic-ai/nomic-embed-text-v1",
475            Self::NomicEmbedTextV15 => "nomic-ai/nomic-embed-text-v1.5",
476            Self::NomicEmbedTextV15Q => "nomic-ai/nomic-embed-text-v1.5",
477            Self::ParaphraseMLMiniLML12V2 => {
478                "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
479            }
480            Self::ParaphraseMLMiniLML12V2Q => "Xenova/paraphrase-multilingual-MiniLM-L12-v2",
481            Self::ParaphraseMLMpnetBaseV2 => {
482                "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
483            }
484            Self::BGESmallZHV15 => "BAAI/bge-small-zh-v1.5",
485            Self::BGELargeZHV15 => "BAAI/bge-large-zh-v1.5",
486            Self::BGEM3 => "BAAI/bge-m3",
487            Self::ModernBertEmbedLarge => "nomic-ai/modernbert-embed-large",
488            Self::MultilingualE5Small => "intfloat/multilingual-e5-small",
489            Self::MultilingualE5Base => "intfloat/multilingual-e5-base",
490            Self::MultilingualE5Large => "intfloat/multilingual-e5-large",
491            Self::MxbaiEmbedLargeV1 => "mixedbread-ai/mxbai-embed-large-v1",
492            Self::MxbaiEmbedLargeV1Q => "mixedbread-ai/mxbai-embed-large-v1",
493            Self::GTEBaseENV15 => "Alibaba-NLP/gte-base-en-v1.5",
494            Self::GTEBaseENV15Q => "Qdrant/gte-base-en-v1.5-onnx-Q",
495            Self::GTELargeENV15 => "Alibaba-NLP/gte-large-en-v1.5",
496            Self::GTELargeENV15Q => "Qdrant/gte-large-en-v1.5-onnx-Q",
497            Self::ClipVitB32 => "openai/clip-vit-base-patch32",
498            Self::JinaEmbeddingsV2BaseCode => "jinaai/jina-embeddings-v2-base-code",
499            Self::JinaEmbeddingsV2BaseEN => "jinaai/jina-embeddings-v2-base-en",
500            Self::EmbeddingGemma300M => "google/embedding-gemma-300M",
501            Self::SnowflakeArcticEmbedXS => "Snowflake/snowflake-arctic-embed-xs",
502            Self::SnowflakeArcticEmbedXSQ => "Snowflake/snowflake-arctic-embed-xs",
503            Self::SnowflakeArcticEmbedS => "Snowflake/snowflake-arctic-embed-s",
504            Self::SnowflakeArcticEmbedSQ => "Snowflake/snowflake-arctic-embed-s",
505            Self::SnowflakeArcticEmbedM => "Snowflake/snowflake-arctic-embed-m",
506            Self::SnowflakeArcticEmbedMQ => "Snowflake/snowflake-arctic-embed-m",
507            Self::SnowflakeArcticEmbedMLong => "Snowflake/snowflake-arctic-embed-m-long",
508            Self::SnowflakeArcticEmbedMLongQ => "Snowflake/snowflake-arctic-embed-m-long",
509            Self::SnowflakeArcticEmbedL => "Snowflake/snowflake-arctic-embed-l",
510            Self::SnowflakeArcticEmbedLQ => "Snowflake/snowflake-arctic-embed-l",
511        }
512    }
513
514    /// Maximum token context per model.
515    pub const fn max_seq_length(self) -> usize {
516        match self {
517            // 256 tokens
518            Self::AllMiniLML6V2
519            | Self::AllMiniLML6V2Q
520            | Self::AllMiniLML12V2
521            | Self::AllMiniLML12V2Q => 256,
522            // 384 tokens
523            Self::AllMpnetBaseV2 => 384,
524            // 8192 tokens
525            Self::BGEM3
526            | Self::NomicEmbedTextV1
527            | Self::NomicEmbedTextV15
528            | Self::NomicEmbedTextV15Q
529            | Self::JinaEmbeddingsV2BaseCode
530            | Self::JinaEmbeddingsV2BaseEN
531            | Self::EmbeddingGemma300M
532            | Self::SnowflakeArcticEmbedMLong
533            | Self::SnowflakeArcticEmbedMLongQ => 8192,
534            // 512 tokens (default)
535            Self::BGESmallENV15
536            | Self::BGESmallENV15Q
537            | Self::BGEBaseENV15
538            | Self::BGEBaseENV15Q
539            | Self::BGELargeENV15
540            | Self::BGELargeENV15Q
541            | Self::BGESmallZHV15
542            | Self::BGELargeZHV15
543            | Self::ParaphraseMLMiniLML12V2
544            | Self::ParaphraseMLMiniLML12V2Q
545            | Self::ParaphraseMLMpnetBaseV2
546            | Self::ModernBertEmbedLarge
547            | Self::MultilingualE5Small
548            | Self::MultilingualE5Base
549            | Self::MultilingualE5Large
550            | Self::MxbaiEmbedLargeV1
551            | Self::MxbaiEmbedLargeV1Q
552            | Self::GTEBaseENV15
553            | Self::GTEBaseENV15Q
554            | Self::GTELargeENV15
555            | Self::GTELargeENV15Q
556            | Self::ClipVitB32
557            | Self::SnowflakeArcticEmbedXS
558            | Self::SnowflakeArcticEmbedXSQ
559            | Self::SnowflakeArcticEmbedS
560            | Self::SnowflakeArcticEmbedSQ
561            | Self::SnowflakeArcticEmbedM
562            | Self::SnowflakeArcticEmbedMQ
563            | Self::SnowflakeArcticEmbedL
564            | Self::SnowflakeArcticEmbedLQ => 512,
565        }
566    }
567
568    /// Whether this is a quantized model (Q suffix).
569    pub const fn is_quantized(self) -> bool {
570        matches!(
571            self,
572            Self::AllMiniLML6V2Q
573                | Self::AllMiniLML12V2Q
574                | Self::BGEBaseENV15Q
575                | Self::BGELargeENV15Q
576                | Self::BGESmallENV15Q
577                | Self::NomicEmbedTextV15Q
578                | Self::ParaphraseMLMiniLML12V2Q
579                | Self::MxbaiEmbedLargeV1Q
580                | Self::GTEBaseENV15Q
581                | Self::GTELargeENV15Q
582                | Self::SnowflakeArcticEmbedXSQ
583                | Self::SnowflakeArcticEmbedSQ
584                | Self::SnowflakeArcticEmbedMQ
585                | Self::SnowflakeArcticEmbedMLongQ
586                | Self::SnowflakeArcticEmbedLQ
587        )
588    }
589
590    /// Whether this model handles image inputs (CLIP).
591    pub const fn is_image_model(self) -> bool {
592        matches!(self, Self::ClipVitB32)
593    }
594
595    /// HuggingFace model code used by fastembed-rs for downloads and caching
596    /// (e.g. `"Xenova/bge-small-en-v1.5"`).
597    ///
598    /// This is the **actual** repository identifier that `hf-hub` uses when
599    /// downloading model weights. Differs from [`model_id`](Self::model_id)
600    /// when the ONNX conversion lives in a separate repo (e.g. `Qdrant/…`,
601    /// `Xenova/…`, `onnx-community/…`).
602    pub const fn model_code(self) -> &'static str {
603        match self {
604            Self::AllMiniLML6V2 => "Qdrant/all-MiniLM-L6-v2-onnx",
605            Self::AllMiniLML6V2Q => "Xenova/all-MiniLM-L6-v2",
606            Self::AllMiniLML12V2 => "Xenova/all-MiniLM-L12-v2",
607            Self::AllMiniLML12V2Q => "Xenova/all-MiniLM-L12-v2",
608            Self::AllMpnetBaseV2 => "Xenova/all-mpnet-base-v2",
609            Self::BGEBaseENV15 => "Xenova/bge-base-en-v1.5",
610            Self::BGEBaseENV15Q => "Qdrant/bge-base-en-v1.5-onnx-Q",
611            Self::BGELargeENV15 => "Xenova/bge-large-en-v1.5",
612            Self::BGELargeENV15Q => "Qdrant/bge-large-en-v1.5-onnx-Q",
613            Self::BGESmallENV15 => "Xenova/bge-small-en-v1.5",
614            Self::BGESmallENV15Q => "Qdrant/bge-small-en-v1.5-onnx-Q",
615            Self::NomicEmbedTextV1 => "nomic-ai/nomic-embed-text-v1",
616            Self::NomicEmbedTextV15 => "nomic-ai/nomic-embed-text-v1.5",
617            Self::NomicEmbedTextV15Q => "nomic-ai/nomic-embed-text-v1.5",
618            Self::ParaphraseMLMiniLML12V2 => "Xenova/paraphrase-multilingual-MiniLM-L12-v2",
619            Self::ParaphraseMLMiniLML12V2Q => "Qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
620            Self::ParaphraseMLMpnetBaseV2 => "Xenova/paraphrase-multilingual-mpnet-base-v2",
621            Self::BGESmallZHV15 => "Xenova/bge-small-zh-v1.5",
622            Self::BGELargeZHV15 => "Xenova/bge-large-zh-v1.5",
623            Self::BGEM3 => "BAAI/bge-m3",
624            Self::ModernBertEmbedLarge => "lightonai/modernbert-embed-large",
625            Self::MultilingualE5Small => "intfloat/multilingual-e5-small",
626            Self::MultilingualE5Base => "intfloat/multilingual-e5-base",
627            Self::MultilingualE5Large => "Qdrant/multilingual-e5-large-onnx",
628            Self::MxbaiEmbedLargeV1 => "mixedbread-ai/mxbai-embed-large-v1",
629            Self::MxbaiEmbedLargeV1Q => "mixedbread-ai/mxbai-embed-large-v1",
630            Self::GTEBaseENV15 => "Alibaba-NLP/gte-base-en-v1.5",
631            Self::GTEBaseENV15Q => "Alibaba-NLP/gte-base-en-v1.5",
632            Self::GTELargeENV15 => "Alibaba-NLP/gte-large-en-v1.5",
633            Self::GTELargeENV15Q => "Alibaba-NLP/gte-large-en-v1.5",
634            Self::ClipVitB32 => "Qdrant/clip-ViT-B-32-text",
635            Self::JinaEmbeddingsV2BaseCode => "jinaai/jina-embeddings-v2-base-code",
636            Self::JinaEmbeddingsV2BaseEN => "jinaai/jina-embeddings-v2-base-en",
637            Self::EmbeddingGemma300M => "onnx-community/embeddinggemma-300m-ONNX",
638            Self::SnowflakeArcticEmbedXS => "snowflake/snowflake-arctic-embed-xs",
639            Self::SnowflakeArcticEmbedXSQ => "snowflake/snowflake-arctic-embed-xs",
640            Self::SnowflakeArcticEmbedS => "snowflake/snowflake-arctic-embed-s",
641            Self::SnowflakeArcticEmbedSQ => "snowflake/snowflake-arctic-embed-s",
642            Self::SnowflakeArcticEmbedM => "Snowflake/snowflake-arctic-embed-m",
643            Self::SnowflakeArcticEmbedMQ => "Snowflake/snowflake-arctic-embed-m",
644            Self::SnowflakeArcticEmbedMLong => "snowflake/snowflake-arctic-embed-m-long",
645            Self::SnowflakeArcticEmbedMLongQ => "snowflake/snowflake-arctic-embed-m-long",
646            Self::SnowflakeArcticEmbedL => "snowflake/snowflake-arctic-embed-l",
647            Self::SnowflakeArcticEmbedLQ => "snowflake/snowflake-arctic-embed-l",
648        }
649    }
650
651    /// String representation matching the enum variant name.
652    pub fn as_str(self) -> &'static str {
653        match self {
654            Self::BGESmallENV15 => "BGESmallENV15",
655            Self::AllMiniLML6V2 => "AllMiniLML6V2",
656            Self::AllMiniLML6V2Q => "AllMiniLML6V2Q",
657            Self::AllMiniLML12V2 => "AllMiniLML12V2",
658            Self::AllMiniLML12V2Q => "AllMiniLML12V2Q",
659            Self::AllMpnetBaseV2 => "AllMpnetBaseV2",
660            Self::BGEBaseENV15 => "BGEBaseENV15",
661            Self::BGEBaseENV15Q => "BGEBaseENV15Q",
662            Self::BGELargeENV15 => "BGELargeENV15",
663            Self::BGELargeENV15Q => "BGELargeENV15Q",
664            Self::BGESmallENV15Q => "BGESmallENV15Q",
665            Self::NomicEmbedTextV1 => "NomicEmbedTextV1",
666            Self::NomicEmbedTextV15 => "NomicEmbedTextV15",
667            Self::NomicEmbedTextV15Q => "NomicEmbedTextV15Q",
668            Self::ParaphraseMLMiniLML12V2 => "ParaphraseMLMiniLML12V2",
669            Self::ParaphraseMLMiniLML12V2Q => "ParaphraseMLMiniLML12V2Q",
670            Self::ParaphraseMLMpnetBaseV2 => "ParaphraseMLMpnetBaseV2",
671            Self::BGESmallZHV15 => "BGESmallZHV15",
672            Self::BGELargeZHV15 => "BGELargeZHV15",
673            Self::BGEM3 => "BGEM3",
674            Self::ModernBertEmbedLarge => "ModernBertEmbedLarge",
675            Self::MultilingualE5Small => "MultilingualE5Small",
676            Self::MultilingualE5Base => "MultilingualE5Base",
677            Self::MultilingualE5Large => "MultilingualE5Large",
678            Self::MxbaiEmbedLargeV1 => "MxbaiEmbedLargeV1",
679            Self::MxbaiEmbedLargeV1Q => "MxbaiEmbedLargeV1Q",
680            Self::GTEBaseENV15 => "GTEBaseENV15",
681            Self::GTEBaseENV15Q => "GTEBaseENV15Q",
682            Self::GTELargeENV15 => "GTELargeENV15",
683            Self::GTELargeENV15Q => "GTELargeENV15Q",
684            Self::ClipVitB32 => "ClipVitB32",
685            Self::JinaEmbeddingsV2BaseCode => "JinaEmbeddingsV2BaseCode",
686            Self::JinaEmbeddingsV2BaseEN => "JinaEmbeddingsV2BaseEN",
687            Self::EmbeddingGemma300M => "EmbeddingGemma300M",
688            Self::SnowflakeArcticEmbedXS => "SnowflakeArcticEmbedXS",
689            Self::SnowflakeArcticEmbedXSQ => "SnowflakeArcticEmbedXSQ",
690            Self::SnowflakeArcticEmbedS => "SnowflakeArcticEmbedS",
691            Self::SnowflakeArcticEmbedSQ => "SnowflakeArcticEmbedSQ",
692            Self::SnowflakeArcticEmbedM => "SnowflakeArcticEmbedM",
693            Self::SnowflakeArcticEmbedMQ => "SnowflakeArcticEmbedMQ",
694            Self::SnowflakeArcticEmbedMLong => "SnowflakeArcticEmbedMLong",
695            Self::SnowflakeArcticEmbedMLongQ => "SnowflakeArcticEmbedMLongQ",
696            Self::SnowflakeArcticEmbedL => "SnowflakeArcticEmbedL",
697            Self::SnowflakeArcticEmbedLQ => "SnowflakeArcticEmbedLQ",
698        }
699    }
700
701    /// Parse a model name (case-insensitive).
702    pub fn parse(s: &str) -> Result<Self, String> {
703        Self::ALL
704            .iter()
705            .find(|m| m.as_str().eq_ignore_ascii_case(s))
706            .copied()
707            .ok_or_else(|| format!("unknown embedding model: {s}"))
708    }
709
710    /// All supported models.
711    pub const ALL: &[Self] = &[
712        Self::BGESmallENV15,
713        Self::AllMiniLML6V2,
714        Self::AllMiniLML6V2Q,
715        Self::AllMiniLML12V2,
716        Self::AllMiniLML12V2Q,
717        Self::AllMpnetBaseV2,
718        Self::BGEBaseENV15,
719        Self::BGEBaseENV15Q,
720        Self::BGELargeENV15,
721        Self::BGELargeENV15Q,
722        Self::BGESmallENV15Q,
723        Self::NomicEmbedTextV1,
724        Self::NomicEmbedTextV15,
725        Self::NomicEmbedTextV15Q,
726        Self::ParaphraseMLMiniLML12V2,
727        Self::ParaphraseMLMiniLML12V2Q,
728        Self::ParaphraseMLMpnetBaseV2,
729        Self::BGESmallZHV15,
730        Self::BGELargeZHV15,
731        Self::BGEM3,
732        Self::ModernBertEmbedLarge,
733        Self::MultilingualE5Small,
734        Self::MultilingualE5Base,
735        Self::MultilingualE5Large,
736        Self::MxbaiEmbedLargeV1,
737        Self::MxbaiEmbedLargeV1Q,
738        Self::GTEBaseENV15,
739        Self::GTEBaseENV15Q,
740        Self::GTELargeENV15,
741        Self::GTELargeENV15Q,
742        Self::ClipVitB32,
743        Self::JinaEmbeddingsV2BaseCode,
744        Self::JinaEmbeddingsV2BaseEN,
745        Self::EmbeddingGemma300M,
746        Self::SnowflakeArcticEmbedXS,
747        Self::SnowflakeArcticEmbedXSQ,
748        Self::SnowflakeArcticEmbedS,
749        Self::SnowflakeArcticEmbedSQ,
750        Self::SnowflakeArcticEmbedM,
751        Self::SnowflakeArcticEmbedMQ,
752        Self::SnowflakeArcticEmbedMLong,
753        Self::SnowflakeArcticEmbedMLongQ,
754        Self::SnowflakeArcticEmbedL,
755        Self::SnowflakeArcticEmbedLQ,
756    ];
757
758    /// Map to `fastembed::EmbeddingModel`.
759    ///
760    /// Only available when an ONNX embedding feature is enabled
761    /// (`embedding-fastembed` or `embedding-fastembed-dynamic-linking`).
762    #[cfg(any(
763        feature = "embedding-fastembed",
764        feature = "embedding-fastembed-dynamic-linking"
765    ))]
766    pub fn as_fastembed(self) -> fastembed::EmbeddingModel {
767        match self {
768            Self::BGESmallENV15 => fastembed::EmbeddingModel::BGESmallENV15,
769            Self::AllMiniLML6V2 => fastembed::EmbeddingModel::AllMiniLML6V2,
770            Self::AllMiniLML6V2Q => fastembed::EmbeddingModel::AllMiniLML6V2Q,
771            Self::AllMiniLML12V2 => fastembed::EmbeddingModel::AllMiniLML12V2,
772            Self::AllMiniLML12V2Q => fastembed::EmbeddingModel::AllMiniLML12V2Q,
773            Self::AllMpnetBaseV2 => fastembed::EmbeddingModel::AllMpnetBaseV2,
774            Self::BGEBaseENV15 => fastembed::EmbeddingModel::BGEBaseENV15,
775            Self::BGEBaseENV15Q => fastembed::EmbeddingModel::BGEBaseENV15Q,
776            Self::BGELargeENV15 => fastembed::EmbeddingModel::BGELargeENV15,
777            Self::BGELargeENV15Q => fastembed::EmbeddingModel::BGELargeENV15Q,
778            Self::BGESmallENV15Q => fastembed::EmbeddingModel::BGESmallENV15Q,
779            Self::NomicEmbedTextV1 => fastembed::EmbeddingModel::NomicEmbedTextV1,
780            Self::NomicEmbedTextV15 => fastembed::EmbeddingModel::NomicEmbedTextV15,
781            Self::NomicEmbedTextV15Q => fastembed::EmbeddingModel::NomicEmbedTextV15Q,
782            Self::ParaphraseMLMiniLML12V2 => fastembed::EmbeddingModel::ParaphraseMLMiniLML12V2,
783            Self::ParaphraseMLMiniLML12V2Q => fastembed::EmbeddingModel::ParaphraseMLMiniLML12V2Q,
784            Self::ParaphraseMLMpnetBaseV2 => fastembed::EmbeddingModel::ParaphraseMLMpnetBaseV2,
785            Self::BGESmallZHV15 => fastembed::EmbeddingModel::BGESmallZHV15,
786            Self::BGELargeZHV15 => fastembed::EmbeddingModel::BGELargeZHV15,
787            Self::BGEM3 => fastembed::EmbeddingModel::BGEM3,
788            Self::ModernBertEmbedLarge => fastembed::EmbeddingModel::ModernBertEmbedLarge,
789            Self::MultilingualE5Small => fastembed::EmbeddingModel::MultilingualE5Small,
790            Self::MultilingualE5Base => fastembed::EmbeddingModel::MultilingualE5Base,
791            Self::MultilingualE5Large => fastembed::EmbeddingModel::MultilingualE5Large,
792            Self::MxbaiEmbedLargeV1 => fastembed::EmbeddingModel::MxbaiEmbedLargeV1,
793            Self::MxbaiEmbedLargeV1Q => fastembed::EmbeddingModel::MxbaiEmbedLargeV1Q,
794            Self::GTEBaseENV15 => fastembed::EmbeddingModel::GTEBaseENV15,
795            Self::GTEBaseENV15Q => fastembed::EmbeddingModel::GTEBaseENV15Q,
796            Self::GTELargeENV15 => fastembed::EmbeddingModel::GTELargeENV15,
797            Self::GTELargeENV15Q => fastembed::EmbeddingModel::GTELargeENV15Q,
798            Self::ClipVitB32 => fastembed::EmbeddingModel::ClipVitB32,
799            Self::JinaEmbeddingsV2BaseCode => fastembed::EmbeddingModel::JinaEmbeddingsV2BaseCode,
800            Self::JinaEmbeddingsV2BaseEN => fastembed::EmbeddingModel::JinaEmbeddingsV2BaseEN,
801            Self::EmbeddingGemma300M => fastembed::EmbeddingModel::EmbeddingGemma300M,
802            Self::SnowflakeArcticEmbedXS => fastembed::EmbeddingModel::SnowflakeArcticEmbedXS,
803            Self::SnowflakeArcticEmbedXSQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedXSQ,
804            Self::SnowflakeArcticEmbedS => fastembed::EmbeddingModel::SnowflakeArcticEmbedS,
805            Self::SnowflakeArcticEmbedSQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedSQ,
806            Self::SnowflakeArcticEmbedM => fastembed::EmbeddingModel::SnowflakeArcticEmbedM,
807            Self::SnowflakeArcticEmbedMQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedMQ,
808            Self::SnowflakeArcticEmbedMLong => fastembed::EmbeddingModel::SnowflakeArcticEmbedMLong,
809            Self::SnowflakeArcticEmbedMLongQ => {
810                fastembed::EmbeddingModel::SnowflakeArcticEmbedMLongQ
811            }
812            Self::SnowflakeArcticEmbedL => fastembed::EmbeddingModel::SnowflakeArcticEmbedL,
813            Self::SnowflakeArcticEmbedLQ => fastembed::EmbeddingModel::SnowflakeArcticEmbedLQ,
814        }
815    }
816}
817
818impl std::fmt::Display for EmbeddingModel {
819    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
820        f.write_str(self.as_str())
821    }
822}
823
824impl std::str::FromStr for EmbeddingModel {
825    type Err = String;
826
827    fn from_str(s: &str) -> Result<Self, Self::Err> {
828        Self::parse(s)
829    }
830}
831
832#[cfg(test)]
833mod tests {
834    use super::*;
835
836    #[test]
837    fn all_count() {
838        assert_eq!(EmbeddingModel::ALL.len(), 44);
839    }
840
841    #[test]
842    fn default_is_bge_small() {
843        assert_eq!(EmbeddingModel::default(), EmbeddingModel::BGESmallENV15);
844    }
845
846    #[test]
847    fn dimension_consistency() {
848        for &m in EmbeddingModel::ALL {
849            let dim = m.dimension();
850            assert!(
851                [384, 512, 768, 1024].contains(&dim),
852                "{m:?}: unexpected dimension {dim}"
853            );
854        }
855    }
856
857    #[test]
858    fn parse_roundtrip() {
859        for &m in EmbeddingModel::ALL {
860            let s = m.as_str();
861            assert_eq!(EmbeddingModel::parse(s).unwrap(), m);
862        }
863    }
864
865    #[test]
866    fn parse_case_insensitive() {
867        assert_eq!(
868            EmbeddingModel::parse("bgesmallenv15").unwrap(),
869            EmbeddingModel::BGESmallENV15
870        );
871        assert_eq!(
872            EmbeddingModel::parse("ALLMINILML6V2").unwrap(),
873            EmbeddingModel::AllMiniLML6V2
874        );
875    }
876
877    #[test]
878    fn parse_unknown_fails() {
879        assert!(EmbeddingModel::parse("NotARealModel").is_err());
880    }
881
882    #[test]
883    fn quantized_flags() {
884        let quantized: Vec<_> = EmbeddingModel::ALL
885            .iter()
886            .filter(|m| m.is_quantized())
887            .copied()
888            .collect();
889        // All Q-suffix variants should be flagged
890        for m in &quantized {
891            assert!(m.as_str().ends_with('Q'), "{m:?} flagged but no Q suffix");
892        }
893        // Non-Q variants should NOT be flagged
894        for &m in EmbeddingModel::ALL {
895            if !m.as_str().ends_with('Q') {
896                assert!(!m.is_quantized(), "{m:?} not Q but flagged quantized");
897            }
898        }
899    }
900
901    #[test]
902    fn image_model_flag() {
903        assert!(EmbeddingModel::ClipVitB32.is_image_model());
904        assert_eq!(
905            EmbeddingModel::ALL
906                .iter()
907                .filter(|m| m.is_image_model())
908                .count(),
909            1
910        );
911    }
912
913    #[test]
914    fn prefix_mapping() {
915        // E5 models have query + doc prefixes
916        for &m in &[
917            EmbeddingModel::MultilingualE5Small,
918            EmbeddingModel::MultilingualE5Base,
919            EmbeddingModel::MultilingualE5Large,
920        ] {
921            assert_eq!(m.query_prefix(), Some("query: "));
922            assert_eq!(m.doc_prefix(), Some("passage: "));
923        }
924        // Snowflake models have query prefix only
925        for &m in &[
926            EmbeddingModel::SnowflakeArcticEmbedXS,
927            EmbeddingModel::SnowflakeArcticEmbedLQ,
928        ] {
929            assert!(m.query_prefix().is_some());
930            assert!(m.doc_prefix().is_none());
931        }
932        // Nomic v1.5 models have query + doc prefixes
933        for &m in &[
934            EmbeddingModel::NomicEmbedTextV15,
935            EmbeddingModel::NomicEmbedTextV15Q,
936        ] {
937            assert_eq!(m.query_prefix(), Some("search_query: "));
938            assert_eq!(m.doc_prefix(), Some("search_document: "));
939        }
940        // Nomic v1 has no prefixes
941        assert!(EmbeddingModel::NomicEmbedTextV1.query_prefix().is_none());
942        assert!(EmbeddingModel::NomicEmbedTextV1.doc_prefix().is_none());
943        // BGE-en-v1.5 + mxbai are asymmetric (query prefix, no doc prefix)
944        for &m in &[
945            EmbeddingModel::BGESmallENV15,
946            EmbeddingModel::MxbaiEmbedLargeV1,
947        ] {
948            assert!(m.query_prefix().is_some());
949            assert!(m.doc_prefix().is_none());
950        }
951        // MiniLM (symmetric) has no prefixes
952        assert!(EmbeddingModel::AllMiniLML6V2.query_prefix().is_none());
953        assert!(EmbeddingModel::AllMiniLML6V2.doc_prefix().is_none());
954    }
955
956    #[test]
957    fn from_str_trait() {
958        let m: EmbeddingModel = "BGESmallENV15".parse().unwrap();
959        assert_eq!(m, EmbeddingModel::BGESmallENV15);
960    }
961
962    #[test]
963    fn display_trait() {
964        assert_eq!(EmbeddingModel::BGESmallENV15.to_string(), "BGESmallENV15");
965    }
966
967    #[test]
968    fn metadata_nonzero() {
969        for &m in EmbeddingModel::ALL {
970            assert!(m.size_mb() > 0, "{m:?}: size_mb is zero");
971            assert!(!m.model_id().is_empty(), "{m:?}: model_id is empty");
972            assert!(!m.model_code().is_empty(), "{m:?}: model_code is empty");
973            assert!(m.max_seq_length() > 0, "{m:?}: max_seq_length is zero");
974        }
975    }
976
977    #[test]
978    fn max_seq_length_values() {
979        assert_eq!(EmbeddingModel::AllMiniLML6V2.max_seq_length(), 256);
980        assert_eq!(EmbeddingModel::AllMpnetBaseV2.max_seq_length(), 384);
981        assert_eq!(EmbeddingModel::BGEM3.max_seq_length(), 8192);
982        assert_eq!(EmbeddingModel::BGESmallENV15.max_seq_length(), 512);
983    }
984
985    /// The MLX-supported set is 13 base models (21 variants counting the `*Q`
986    /// aliases, which share weights with their non-quantized twin). Locked down
987    /// so adding a variant to `mlx_supported` without verifying its architecture
988    /// against the real weight repo fails here.
989    #[test]
990    fn mlx_supported_set_is_exact() {
991        let supported: Vec<&str> = EmbeddingModel::ALL
992            .iter()
993            .filter(|m| m.mlx_supported())
994            .map(|m| m.as_str())
995            .collect();
996        assert_eq!(
997            supported,
998            vec![
999                "BGESmallENV15",
1000                "AllMiniLML6V2",
1001                "AllMiniLML6V2Q",
1002                "AllMiniLML12V2",
1003                "AllMiniLML12V2Q",
1004                "BGEBaseENV15",
1005                "BGEBaseENV15Q",
1006                "BGELargeENV15",
1007                "BGELargeENV15Q",
1008                "BGESmallENV15Q",
1009                "ParaphraseMLMiniLML12V2",
1010                "ParaphraseMLMiniLML12V2Q",
1011                "BGESmallZHV15",
1012                "MultilingualE5Small",
1013                "MxbaiEmbedLargeV1",
1014                "MxbaiEmbedLargeV1Q",
1015                "SnowflakeArcticEmbedXS",
1016                "SnowflakeArcticEmbedXSQ",
1017                "SnowflakeArcticEmbedS",
1018                "SnowflakeArcticEmbedSQ",
1019                "SnowflakeArcticEmbedM",
1020                "SnowflakeArcticEmbedMQ",
1021                "SnowflakeArcticEmbedL",
1022                "SnowflakeArcticEmbedLQ",
1023            ]
1024        );
1025    }
1026
1027    /// Architectures known not to be loadable by the MLX vanilla-BERT forward
1028    /// pass must stay excluded — regression guard for the arctic-m-long case,
1029    /// where a NomicBert model was listed as supported and failed at load.
1030    #[test]
1031    fn mlx_excludes_non_vanilla_bert() {
1032        for &m in &[
1033            EmbeddingModel::SnowflakeArcticEmbedMLong, // NomicBertModel
1034            EmbeddingModel::SnowflakeArcticEmbedMLongQ,
1035            EmbeddingModel::NomicEmbedTextV1,
1036            EmbeddingModel::NomicEmbedTextV15,
1037            EmbeddingModel::GTEBaseENV15,            // NewModel
1038            EmbeddingModel::GTELargeENV15,           // NewModel
1039            EmbeddingModel::MultilingualE5Base,      // XLM-R
1040            EmbeddingModel::MultilingualE5Large,     // XLM-R
1041            EmbeddingModel::ParaphraseMLMpnetBaseV2, // XLM-R
1042            EmbeddingModel::BGEM3,                   // XLM-R
1043            EmbeddingModel::AllMpnetBaseV2,          // MPNet
1044            EmbeddingModel::JinaEmbeddingsV2BaseEN,  // JinaBert
1045            EmbeddingModel::JinaEmbeddingsV2BaseCode,
1046            EmbeddingModel::ModernBertEmbedLarge, // ModernBERT
1047            EmbeddingModel::EmbeddingGemma300M,
1048            EmbeddingModel::ClipVitB32,
1049            EmbeddingModel::BGELargeZHV15, // BertModel but no safetensors
1050        ] {
1051            assert!(!m.mlx_supported(), "{m:?} must not be MLX-supported");
1052        }
1053    }
1054}
llm_kernel/embedding/catalog.rs

llm_kernel/embedding/
catalog.rs
