EmbeddingGemma 2 — Text Only — BF16 — MLX

A modular deployment export of Google DeepMind's EmbeddingGemma 2. Only the text backbone is retained; the audio encoder and audio projection are absent, as are the vision encoder and projection. No training, distillation, or quantization was applied.

Retained weights preserve the upstream BF16 precision.

Weight size is not total runtime memory. 270M/440M in repository names are rounded deployment sizes.

Install

pip install "mlx>=0.32.3" "transformers>=5.19.0"
pip install "git+https://github.com/Blaizzy/mlx-vlm.git@3d87e88402f307efbf68e568971aa887ee7d9ed0"

EmbeddingGemma 2 currently requires the pinned MLX-VLM implementation revision above. The MLX package was tested with MLX 0.32.3 and Transformers 5.19.0.

Text and code

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("jayyun98/embeddinggemma-2-text-270m-mlx-bf16")
model = load_embedding_model(path)
tokenizer = AutoTokenizer.from_pretrained(path)
inputs = tokenizer([
    "task: search result | query: What causes the northern lights?",
    "title: none | text: Charged particles from the sun cause the northern lights.",
], padding=True, return_tensors="np")
vectors = model(**{key: mx.array(value) for key, value in inputs.items()}).text_embeds
mx.eval(vectors)
print(vectors[:1] @ vectors[1:].T)

# Optional MRL truncation. Normalize after slicing.
vectors = vectors[:, :256]
vectors /= mx.linalg.norm(vectors, axis=-1, keepdims=True)

Use SearchQuery for search queries, CodeRetrieval for code-search queries, and Document for corpus items. For the MLX API, prepend the corresponding literal prefix from config_sentence_transformers.json, as shown above. For titled documents, use title: {title} | text: {content} without another prefix. Use matching dimensions for queries and documents, and normalize after truncation. Do not cast this model to FP16; use BF16 or FP32 inference.

Verification

Strict native MLX loading and Apple Silicon GPU inference were tested. Cross-runtime arithmetic and BF16 rounding can produce small differences from PyTorch FP32. Checks cover English/Korean search and document text, code queries, and 128/256/512-dimensional normalized vectors.

Measurements compare with the pinned full Google checkpoint. They are small numerical and loading checks, not MTEB results, a retrieval-quality evaluation, or a speed benchmark. The original prompts, pooling behavior, tokenizer and processor assets are retained. Processor metadata does not restore the removed encoder weights. Image, audio and video encoders are unavailable.

• conversion.json: storage precision, removal and conversion checks.
• verification.json: runtime, numerical results and tested environment.
• Original source revision: 914f7f89142e33e77833254d9c9b90c3cef7303b.
• MLX-VLM implementation: 3d87e884.

Attribution and license

Original weights and tokenizer/processor assets: Google DeepMind. MLX implementation: MLX-VLM contributors. This is an independent derivative deployment export, not an official Google release. The Apache 2.0 LICENSE and NOTICE are included. See the original model card for training, intended use and limitations.