Add EmbeddingGemma 2 MLX build

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

1_Pooling/config.json

@@ -0,0 +1,5 @@
+{
+    "embedding_dimension": 768,
+    "pooling_mode": "mean",
+    "include_prompt": true
+}

2_Normalize/config.json

@@ -0,0 +1,4 @@
+{
+    "module_input_name": "sentence_embedding",
+    "module_output_name": "sentence_embedding"
+}

README.md

@@ -0,0 +1,82 @@
+---
+base_model: google/embeddinggemma-2
+library_name: mlx
+pipeline_tag: sentence-similarity
+license: apache-2.0
+tags:
+- mlx
+- embeddings
+- sentence-similarity
+- quantized
+- krill
+- 8bit
+---
+
+# EmbeddingGemma 2, MLX 8bit
+
+A quantized MLX build of [google/embeddinggemma-2](https://huggingface.co/google/embeddinggemma-2) (740M parameters; text, image, audio and video embeddings in one `model.safetensors`), made with [Krill](https://github.com/srvsngh99/Krill). Bits: **8**. Weights file: **806.2 MB** (bf16 original: 1488.9 MB).
+
+## What this build is
+
+Affine 8-bit quantization, group size 64, applied to every quantizable tensor.
+
+Produced with:
+
+```
+krill quantize <bf16-dir> --bits 8 --group-size 64 --dtype bf16 --output-dir <out>
+```
+
+Gate result: **PASS** (8-bit class). Drops against bf16: SciFact -0.03, Hindi -0.03, image R@1 0.0 points (negative means the build scored higher). Gate: SciFact and Hindi drop <= 1.0 nDCG point, image R@1 drop <= 1.0 point (variants: mixed-mxfp8, 8bit, 6bit, 6bit-dyn).
+
+## Quality against bf16
+
+Cosine similarity is measured against Google's fp32 sentence-transformers reference vectors; retrieval is nDCG@10 or Recall@k. The bf16 column is the unmodified Google weights run through Krill.
+
+| metric | bf16 | this build |
+|---|---|---|
+| Text cosine to Google fp32 (min / mean) | 1.000 / 1.000 | 1.000 / 1.000 |
+| Image cosine (min / mean) | 0.999 / 1.000 | 0.999 / 1.000 |
+| Mixed cosine (min / mean) | 1.000 / 1.000 | 0.999 / 1.000 |
+| Audio cosine (min / mean) | 1.000 / 1.000 | 1.000 / 1.000 |
+| Video cosine (min / mean) | 0.999 / 1.000 | 0.999 / 1.000 |
+| SciFact nDCG@10, 768 d | 86.92 | 86.95 |
+| SciFact nDCG@10, 256 d | 84.44 | 84.32 |
+| Hindi nDCG@10 (IndicQARetrieval) | 72.78 | 72.81 |
+| Kannada nDCG@10 (IndicQARetrieval) | 73.25 | 73.25 |
+| Flickr image-to-text R@1 / R@5 | 99.5 / 100.0 | 99.5 / 100.0 |
+| Flickr text-to-image R@1 / R@5 | 96.1 / 99.7 | 96.1 / 99.7 |
+| Clotho audio-to-text R@1 / R@5 | 25.0 / 61.7 | 23.3 / 58.3 |
+| Clotho text-to-audio R@1 / R@5 | 30.0 / 58.3 | 30.0 / 55.0 |
+| Ranking agreement, Spearman raw / Document | 1.000 / 1.000 | 1.000 / 1.000 |
+| Top-1 neighbour unchanged, raw / Document (of 21) | 21 / 21 | 21 / 20 |
+
+SciFact is MTEB `mteb/scifact` (300 queries, 5,183 documents). Hindi and Kannada are the `hi` and `kn` splits of MTEB `IndicQARetrieval`; the corpora are small, so treat differences under about 0.5 point as noise. Flickr is the first 200 images of the Flickr30k 1K test split. Clotho is 60 clips of the Clotho v2 test split (one clip is 1.7 points of Recall@1). Test sets are small; "PASS" means no measurable loss on these sets, not "equal".
+
+## Speed and memory (Apple M4 Pro)
+
+| contender | weights MB | single query p50 ms | batch-32 ~256 tok docs/s | batch-32 ~1k tok docs/s | peak memory MB |
+|---|---|---|---|---|---|
+| Krill bf16 (fp32 compute) | 1488.9 | 10.2 | 39 | 9 | 3234 |
+| Krill 8bit (this build) | 806.2 | 8.8 | 37 | 8 | 2498 |
+| Google sentence-transformers, MPS bf16 | 1488.9 | 26.2 | 26 | 5 | 17283 |
+| Ollama `embeddinggemma-2` (nvfp4, 1.3 GB with towers) | 1300 | 25.2 | 32 | 8 | 1476 |
+| Unsloth Q8_0 GGUF (text only, llama.cpp Metal) | 309.9 | 16.4 | 33 | 7 | 1825 |
+
+Unsloth and Ollama are text only; Krill rows include the vision and audio towers in the weights. Image encoding for this build: 382 ms p50 for one 640x480 PNG. Cold first request: 1086 ms. Single-query p95: 12.9 ms.
+
+Method: Apple M4 Pro, 24 GB. One contender at a time, with the server started and stopped for every repetition, 10 discarded warm-up requests, and 3 repetitions per metric reported as the median. Before each repetition the harness required 1-minute load below 3, at least 50% free memory, no thermal limit and AC power. A contender counts as stable when the spread (max - min over the median) of every metric across its 3 repetitions is at most 10%; otherwise it was re-measured. Krill rows use fp32 compute (the default). Full protocol: `docs/bench/embeddinggemma2-2026-10-08.md` in the Krill repository.
+
+## Usage with Krill
+
+```
+krill pull embeddinggemma-2-8bit
+krill serve
+curl localhost:57455/v1/embeddings -H 'Content-Type: application/json' \
+  -H "Authorization: Bearer $KRILL_API_KEY" -d '{"model":"embeddinggemma-2-8bit","input":"What causes the northern lights?","task":"SearchQuery"}'
+```
+
+Use `"task": "SearchQuery"` for queries and `"Document"` for documents. Dimensions: 768 (default), 512, 256, 128 (Matryoshka). Context: 8,192 tokens. Compute dtype: float32 (default) or bfloat16; float16 is not supported.
+
+## Other builds
+
+`mixed-mxfp8`, `8bit`, `6bit`, `6bit-dyn`, `5bit`, `4bit-g32`, `4bit-dyn`, `4bit-dyn-text` and `nvfp4`, all under `srv-sngh/embeddinggemma-2-mlx-<name>`.

chat_template.jinja

@@ -0,0 +1 @@
+{%- for msg in messages if msg.get('role') == 'system' -%}{%- if msg.get('content') is string -%}{{ msg['content'] }}{%- else -%}{%- for item in msg['content'] if item.get('type') == 'text' -%}{{ item['text'] }}{%- endfor -%}{%- endif -%}{%- endfor -%}{%- for msg in messages if msg.get('role') != 'system' -%}{%- if msg.get('content') is string -%}{{ msg['content'] }}{%- else -%}{%- set existing_text = msg['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join -%}{%- set has_manual_placeholders = ('<|image|>' in existing_text) or ('<|video|>' in existing_text) or ('<|audio|>' in existing_text) -%}{%- for item in msg['content'] -%}{%- if item.get('type') == 'text' -%}{{ item['text'] }}{%- elif not has_manual_placeholders and item.get('type') == 'image' -%}<|image|>{%- elif not has_manual_placeholders and item.get('type') == 'video' -%}<|video|>{%- elif not has_manual_placeholders and item.get('type') == 'audio' -%}<|audio|>{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}

config.json

@@ -0,0 +1,185 @@
+{
+  "architectures" : [
+    "EmbeddingGemma2Model"
+  ],
+  "audio_config" : {
+    "_name_or_path" : "",
+    "architectures" : null,
+    "attention_chunk_size" : 12,
+    "attention_context_left" : 13,
+    "attention_context_right" : 0,
+    "attention_invalid_logits_value" : -1000000000,
+    "attention_logit_cap" : 50,
+    "chunk_size_feed_forward" : 0,
+    "conv_kernel_size" : 5,
+    "dtype" : "bfloat16",
+    "gradient_clipping" : 10000000000,
+    "hidden_act" : "silu",
+    "hidden_size" : 1024,
+    "id2label" : {
+      "0" : "LABEL_0",
+      "1" : "LABEL_1"
+    },
+    "initializer_range" : 0.02,
+    "is_encoder_decoder" : false,
+    "label2id" : {
+      "LABEL_0" : 0,
+      "LABEL_1" : 1
+    },
+    "model_type" : "gemma4_audio",
+    "num_attention_heads" : 8,
+    "num_hidden_layers" : 12,
+    "output_attentions" : false,
+    "output_hidden_states" : false,
+    "output_proj_dims" : 1536,
+    "problem_type" : null,
+    "residual_weight" : 0.5,
+    "return_dict" : true,
+    "rms_norm_eps" : 9.9999999999999995e-07,
+    "subsampling_conv_channels" : [
+      128,
+      32
+    ],
+    "use_clipped_linears" : true
+  },
+  "audio_token_id" : 258881,
+  "boa_token_id" : 256000,
+  "boi_token_id" : 255999,
+  "dtype" : "bfloat16",
+  "eoa_token_index" : 258883,
+  "eoi_token_id" : 258882,
+  "image_token_id" : 258880,
+  "initializer_range" : 0.02,
+  "model_type" : "embedding_gemma2",
+  "quantization" : {
+    "bits" : 8,
+    "group_size" : 64
+  },
+  "quantization_config" : {
+    "bits" : 8,
+    "group_size" : 64
+  },
+  "text_config" : {
+    "attention_bias" : false,
+    "attention_dropout" : 0,
+    "bos_token_id" : 2,
+    "dtype" : "bfloat16",
+    "embedding_dim" : 768,
+    "eos_token_id" : 1,
+    "head_dim" : 256,
+    "hidden_activation" : "gelu_pytorch_tanh",
+    "hidden_size" : 512,
+    "hidden_size_per_layer_input" : 512,
+    "initializer_range" : 0.02,
+    "intermediate_size" : 2048,
+    "layer_types" : [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings" : 262144,
+    "model_type" : "embedding_gemma2_text",
+    "num_attention_heads" : 4,
+    "num_hidden_layers" : 24,
+    "num_key_value_heads" : 2,
+    "pad_token_id" : 0,
+    "per_layer_config" : {
+      "05" : {
+        "head_dim" : 512,
+        "num_key_value_heads" : 1
+      },
+      "11" : {
+        "head_dim" : 512,
+        "num_key_value_heads" : 1
+      },
+      "17" : {
+        "head_dim" : 512,
+        "num_key_value_heads" : 1
+      },
+      "23" : {
+        "head_dim" : 512,
+        "num_key_value_heads" : 1
+      }
+    },
+    "rms_norm_eps" : 9.9999999999999995e-07,
+    "rope_parameters" : {
+      "full_attention" : {
+        "rope_theta" : 1000000,
+        "rope_type" : "default"
+      },
+      "sliding_attention" : {
+        "rope_theta" : 10000,
+        "rope_type" : "default"
+      }
+    },
+    "sliding_window" : 512,
+    "vocab_size" : 262144
+  },
+  "transformers_version" : "5.18.0.dev0",
+  "video_token_id" : 258884,
+  "vision_config" : {
+    "_name_or_path" : "",
+    "architectures" : null,
+    "attention_bias" : false,
+    "attention_dropout" : 0,
+    "chunk_size_feed_forward" : 0,
+    "default_output_length" : 280,
+    "dtype" : "bfloat16",
+    "global_head_dim" : 64,
+    "head_dim" : 64,
+    "hidden_activation" : "gelu_pytorch_tanh",
+    "hidden_size" : 768,
+    "id2label" : {
+      "0" : "LABEL_0",
+      "1" : "LABEL_1"
+    },
+    "initializer_range" : 0.02,
+    "intermediate_size" : 3072,
+    "is_encoder_decoder" : false,
+    "label2id" : {
+      "LABEL_0" : 0,
+      "LABEL_1" : 1
+    },
+    "max_position_embeddings" : 131072,
+    "model_type" : "gemma4_vision",
+    "num_attention_heads" : 12,
+    "num_hidden_layers" : 16,
+    "num_key_value_heads" : 12,
+    "output_attentions" : false,
+    "output_hidden_states" : false,
+    "patch_size" : 16,
+    "pooling_kernel_size" : 3,
+    "position_embedding_size" : 10240,
+    "problem_type" : null,
+    "return_dict" : true,
+    "rms_norm_eps" : 9.9999999999999995e-07,
+    "rope_parameters" : {
+      "rope_theta" : 100,
+      "rope_type" : "axial"
+    },
+    "standardize" : false,
+    "use_clipped_linears" : false
+  },
+  "vision_soft_tokens_per_image" : 280
+}

config_sentence_transformers.json

@@ -0,0 +1,38 @@
+{
+  "__version__": {
+    "pytorch": "2.14.0+cu130",
+    "sentence_transformers": "6.1.0",
+    "transformers": "5.18.0.dev0"
+  },
+  "default_prompt_name": null,
+  "model_type": "SentenceTransformer",
+  "prompts": {
+    "BitextMining": "task: search result | query: ",
+    "Classification": "task: classification | query: ",
+    "Clustering": "task: clustering | query: ",
+    "CodeRetrieval": "task: code retrieval | query: ",
+    "Document": "title: none | text: ",
+    "FactChecking": "task: fact checking | query: ",
+    "InstructionRetrieval": "task: code retrieval | query: ",
+    "MultilabelClassification": "task: classification | query: ",
+    "PairClassification": "task: sentence similarity | query: ",
+    "QuestionAnswering": "task: question answering | query: ",
+    "Reranking": "task: search result | query: ",
+    "Retrieval": "task: search result | query: ",
+    "Retrieval-document": "title: none | text: ",
+    "Retrieval-query": "task: search result | query: ",
+    "STS": "task: sentence similarity | query: ",
+    "SearchQuery": "task: search result | query: ",
+    "SentenceSimilarity": "task: sentence similarity | query: ",
+    "Summarization": "task: sentence similarity | query: ",
+    "document": "title: none | text: ",
+    "query": "task: search result | query: "
+  },
+  "requirements": {
+    "sentence-transformers": {
+      "reason": "Older versions ignore the order of multimodal dict inputs and cannot put several images or texts into a single embedding.",
+      "specifier": ">=6.1.0"
+    }
+  },
+  "similarity_fn_name": "cosine"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:122319f6d385e0072e18dccd26b190b60959b2d6f97a28b62908120166e3622b
+size 806249780

modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.transformer.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.sentence_transformer.modules.pooling.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.base.modules.normalize.Normalize"
+  }
+]

preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+  "dither": 0.0,
+  "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+  "feature_size": 128,
+  "fft_length": 512,
+  "fft_overdrive": false,
+  "frame_length": 320,
+  "hop_length": 160,
+  "input_scale_factor": 1.0,
+  "max_frequency": 8000.0,
+  "mel_floor": 0.001,
+  "min_frequency": 0.0,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "per_bin_mean": null,
+  "per_bin_stddev": null,
+  "preemphasis": 0.0,
+  "preemphasis_htk_flavor": true,
+  "return_attention_mask": true,
+  "sampling_rate": 16000
+}

processor_config.json

@@ -0,0 +1,78 @@
+{
+  "audio_ms_per_token": 40,
+  "audio_seq_length": 280,
+  "feature_extractor": {
+    "dither": 0.0,
+    "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+    "feature_size": 128,
+    "fft_length": 512,
+    "fft_overdrive": false,
+    "frame_length": 320,
+    "hop_length": 160,
+    "input_scale_factor": 1.0,
+    "max_frequency": 8000.0,
+    "mel_floor": 0.001,
+    "min_frequency": 0.0,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "per_bin_mean": null,
+    "per_bin_stddev": null,
+    "preemphasis": 0.0,
+    "preemphasis_htk_flavor": true,
+    "return_attention_mask": true,
+    "sampling_rate": 16000
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_processor_type": "Gemma4ImageProcessor",
+    "image_seq_length": 280,
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 280,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098
+  },
+  "image_seq_length": 280,
+  "processor_class": "EmbeddingGemma2Processor",
+  "video_processor": {
+    "add_timestamps": false,
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "fps": 1,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_frames": 32,
+    "max_soft_tokens": 140,
+    "overflow_strategy": "uniform",
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "video_processor_type": "EmbeddingGemma2VideoProcessor"
+  }
+}

sentence_bert_config.json

@@ -0,0 +1,27 @@
+{
+    "transformer_task": "feature-extraction",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state"
+        },
+        "image": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state"
+        },
+        "audio": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state"
+        },
+        "video": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state"
+        },
+        "message": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state",
+            "format": "structured"
+        }
+    },
+    "module_output_name": "token_embeddings"
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d777ef5bdc1aa36227abdfb77c3e49e7b9c892d16e1b6bda41c393504828be4
+size 32170510

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e594c8a90eb08d8bda498ff4747977dc827ae0c3c56b5c0d41a605a22d02ef03
+size 4689013

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>",
+    "video_token": "<|video|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "right",
+  "processor_class": "EmbeddingGemma2Processor",
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>",
+  "video_token": "<|video|>"
+}