Add 8bit MLX conversion of Google EmbeddingGemma 2

.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,94 @@
+---
+base_model: google/embeddinggemma-2
+language:
+- multilingual
+library_name: mlx
+license: apache-2.0
+pipeline_tag: feature-extraction
+tags:
+- mlx
+- mlx-vlm
+- embedding
+- feature-extraction
+- sentence-similarity
+- multimodal
+- image-feature-extraction
+- audio-feature-extraction
+- video-feature-extraction
+---
+# nativ-community/embeddinggemma-2-8bit
+
+MLX conversion of [Google EmbeddingGemma 2](https://huggingface.co/google/embeddinggemma-2) in **8-bit affine** format, retaining text, image, audio, and video encoders. The model produces normalized 768-dimensional embeddings. See the [original model card](https://huggingface.co/google/embeddinggemma-2) for intended use, training information, evaluation, and limitations.
+
+Quantization mode: `affine`, 8 bits, group size 64. The standard MLX-VLM policy quantizes the text encoder and audio projection. The vision and audio towers and vision projection remain BF16.
+
+Weight storage: 1.234 GB (decimal). Source revision: [`914f7f89142e33e77833254d9c9b90c3cef7303b`](https://huggingface.co/google/embeddinggemma-2/tree/914f7f89142e33e77833254d9c9b90c3cef7303b). Converted with [MLX-VLM revision `3d87e884`](https://github.com/Blaizzy/mlx-vlm/tree/3d87e88402f307efbf68e568971aa887ee7d9ed0), from branch `pc/embeddinggemma-2`, and MLX 0.32.3.
+
+## Text embeddings
+
+Install the implementation revision that includes EmbeddingGemma 2 support:
+
+```sh
+pip install "git+https://github.com/Blaizzy/mlx-vlm.git@3d87e88402f307efbf68e568971aa887ee7d9ed0"
+pip install -U "mlx>=0.32.3" "transformers>=5.18.0"
+```
+
+```python
+import mlx.core as mx
+from transformers import AutoTokenizer
+from mlx_vlm.embedding_loader import load_embedding_model
+from mlx_vlm.utils import get_model_path
+
+path = get_model_path("nativ-community/embeddinggemma-2-8bit")
+model = load_embedding_model(path)
+tokenizer = AutoTokenizer.from_pretrained(path)
+inputs = tokenizer([
+    "task: search result | query: Which planet is known as the Red Planet?",
+    "title: none | text: Mars is known as the Red Planet.",
+    "title: none | text: Venus is often called Earth's twin.",
+], padding=True, return_tensors="np")
+embeddings = model(**{key: mx.array(value) for key, value in inputs.items()}).text_embeds
+print(embeddings[:1] @ embeddings[1:].T)
+
+# Optional Matryoshka truncation: use 128, 256, 512, or 768 dimensions.
+embeddings = embeddings[:, :256]
+embeddings /= mx.linalg.norm(embeddings, axis=-1, keepdims=True)
+```
+
+Apply the appropriate task prefixes from `config_sentence_transformers.json`. Queries and documents must use the same embedding dimension. Keep non-quantized weights and activations in BF16; do not cast this model to float16.
+
+## Multimodal preprocessing
+
+All modality weights and processor configuration files are included. Image, audio, and video preprocessing requires a Transformers build that exposes `EmbeddingGemma2Processor`. Validation used the model's supplied upstream Transformers 5.18.0.dev0 build; stock PyPI Transformers 5.18.0 does not yet expose this processor. The text-only example above uses `AutoTokenizer` and does not require that processor. With a compatible processor build:
+
+```python
+from mlx_vlm import load
+
+model, processor = load("nativ-community/embeddinggemma-2-8bit")
+inputs = processor(images=[image], return_tensors="np")  # image is a PIL image
+embedding = model(**{key: mx.array(value) for key, value in inputs.items()}).text_embeds
+```
+
+## Conversion checks
+
+All checked outputs were finite, unit-normalized, and had 768 dimensions. The text retrieval smoke test ranked the Mars passage above Venus. Checks cover six multilingual text inputs and synthetic image, audio, two-frame video, and text+image inputs; they are numerical smoke checks, not MTEB or a retrieval-quality benchmark. Cosine and absolute errors compare against the original checkpoint running in PyTorch FP32.
+
+| Input | Minimum cosine vs FP32 | Maximum absolute error |
+| --- | ---: | ---: |
+| audio | 0.999804 | 0.002476 |
+| image | 0.999890 | 0.001937 |
+| text | 0.999814 | 0.002299 |
+| text_image | 0.999823 | 0.002105 |
+| video | 0.999836 | 0.002130 |
+
+The full measurements, including truncated-vector comparisons, are in `validation.json`.
+
+## Reproduce conversion
+
+Using the implementation and compatible processor environment described above:
+
+```sh
+python -m mlx_vlm convert --hf-path google/embeddinggemma-2 --revision 914f7f89142e33e77833254d9c9b90c3cef7303b --mlx-path embeddinggemma-2-8bit --dtype bfloat16 -q --q-mode affine --q-bits 8 --q-group-size 64
+```
+
+The source model declares Apache-2.0 licensing. This conversion preserves that license and attribution to Google.

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,187 @@
+{
+    "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.0,
+        "attention_logit_cap": 50.0,
+        "chunk_size_feed_forward": 0,
+        "conv_kernel_size": 5,
+        "dtype": "bfloat16",
+        "gradient_clipping": 10000000000.0,
+        "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": 1e-06,
+        "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": {
+        "group_size": 64,
+        "bits": 8,
+        "mode": "affine"
+    },
+    "quantization_config": {
+        "group_size": 64,
+        "bits": 8,
+        "mode": "affine"
+    },
+    "text_config": {
+        "attention_bias": false,
+        "attention_dropout": 0.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": 1e-06,
+        "rope_parameters": {
+            "full_attention": {
+                "rope_theta": 1000000.0,
+                "rope_type": "default"
+            },
+            "sliding_attention": {
+                "rope_theta": 10000.0,
+                "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.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": 1e-06,
+        "rope_parameters": {
+            "rope_theta": 100.0,
+            "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:6b7e97f9687ad422ab3625a1cdf6002892187d7dd04d01b2a09253fa1e5d4cfd
+size 1234238431

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_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|>"
+}

validation.json

@@ -0,0 +1,77 @@
+{
+  "source_model": "google/embeddinggemma-2",
+  "source_revision": "914f7f89142e33e77833254d9c9b90c3cef7303b",
+  "mlx_vlm_revision": "3d87e88402f307efbf68e568971aa887ee7d9ed0",
+  "variant": "8bit",
+  "hardware": "Apple M3 Ultra",
+  "mlx_version": "0.32.3",
+  "transformers_version": "5.18.0.dev0",
+  "validation_scope": "Six multilingual text inputs and synthetic image, audio, two-frame video, and text+image; numerical smoke checks, not a retrieval benchmark.",
+  "tensor_count": 1816,
+  "tensor_dtypes": [
+    "BF16",
+    "U32"
+  ],
+  "weight_bytes": 1234238431,
+  "cases": {
+    "audio": {
+      "shape": [
+        1,
+        768
+      ],
+      "finite": true,
+      "minimum_cosine_vs_torch_fp32": 0.9998035070404524,
+      "maximum_absolute_error_vs_torch_fp32": 0.002476184321109907,
+      "minimum_cosine_vs_source_bf16": 0.9997744741131873,
+      "minimum_truncated_cosine_vs_torch_fp32": 0.9998247233031456
+    },
+    "image": {
+      "shape": [
+        1,
+        768
+      ],
+      "finite": true,
+      "minimum_cosine_vs_torch_fp32": 0.9998904477000761,
+      "maximum_absolute_error_vs_torch_fp32": 0.0019366481302732774,
+      "minimum_cosine_vs_source_bf16": 0.9999318244275229,
+      "minimum_truncated_cosine_vs_torch_fp32": 0.9998950619484888
+    },
+    "text": {
+      "shape": [
+        6,
+        768
+      ],
+      "finite": true,
+      "minimum_cosine_vs_torch_fp32": 0.9998142728897385,
+      "maximum_absolute_error_vs_torch_fp32": 0.0022990193383245483,
+      "minimum_cosine_vs_source_bf16": 0.9997514709446282,
+      "minimum_truncated_cosine_vs_torch_fp32": 0.9998170484369469,
+      "retrieval_scores_venus_mars": [
+        0.678564727306366,
+        0.8342599868774414
+      ]
+    },
+    "text_image": {
+      "shape": [
+        1,
+        768
+      ],
+      "finite": true,
+      "minimum_cosine_vs_torch_fp32": 0.9998230728231623,
+      "maximum_absolute_error_vs_torch_fp32": 0.002104893632396264,
+      "minimum_cosine_vs_source_bf16": 0.999874734632146,
+      "minimum_truncated_cosine_vs_torch_fp32": 0.999830769497079
+    },
+    "video": {
+      "shape": [
+        1,
+        768
+      ],
+      "finite": true,
+      "minimum_cosine_vs_torch_fp32": 0.9998359946491495,
+      "maximum_absolute_error_vs_torch_fp32": 0.0021304845535331423,
+      "minimum_cosine_vs_source_bf16": 0.9998825464984411,
+      "minimum_truncated_cosine_vs_torch_fp32": 0.9998420387979006
+    }
+  }
+}