EmbeddingGemma 2 text encoder on Core ML (Neural Engine): 7 shared-weight functions, bf16 token table

.gitattributes

@@ -33,3 +33,5 @@ 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
+embeddings.bf16 filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

EmbeddingGemma2Text.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

EmbeddingGemma2Text.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

EmbeddingGemma2Text.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "0EA7802A-D153-497D-B221-C103BA870AB7": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "DE6B4D44-991D-42A8-8048-7490E417887D": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "0EA7802A-D153-497D-B221-C103BA870AB7"
+}

README.md

@@ -0,0 +1,81 @@
+---
+license: apache-2.0
+base_model: google/embeddinggemma-2
+library_name: coreml
+pipeline_tag: feature-extraction
+language:
+- multilingual
+tags:
+- coreml
+- apple-neural-engine
+- on-device
+- embeddings
+- sentence-similarity
+- gemma
+---
+
+# EmbeddingGemma 2 (text) — Core ML, Neural Engine
+
+The text encoder of [google/embeddinggemma-2](https://huggingface.co/google/embeddinggemma-2) converted to Core ML so it
+runs entirely on the Apple Neural Engine. Weights are the original ones (stored fp16 in the model; the token table is
+the original bf16). The vision and audio encoders are not included.
+
+Output: the same 768-d, L2-normalized embedding as `SentenceTransformer("google/embeddinggemma-2").encode(...)`
+(mean pooling over tokens, then normalization). Matryoshka truncation to 512/256/128 works as in the original: keep
+the first N values and re-normalize.
+
+## Files
+
+| File | What |
+|---|---|
+| `EmbeddingGemma2Text.mlpackage` | ML Program, fp16, 7 functions sharing one set of weights (271 MB) |
+| `embeddings.bf16` | token embedding table, 262,144 × 512 raw bfloat16 (256 MB); look up rows, multiply by √512 |
+| `tokenizer.json` | the original Gemma tokenizer, unchanged |
+| `config.json` | shapes, function list, token ids, task prefixes |
+
+## Functions
+
+Every function has fixed shapes, which is what keeps it on the Neural Engine. Variable-length (enumerated) shapes
+put the whole graph on the CPU.
+
+| Function | Inputs | Output |
+|---|---|---|
+| `embed_32` … `embed_512` | `inputs_embeds` [1, S, 512] fp16, `attention_mask` [1, S] fp16 (1 = token, right padding) | `embedding` [1, 768] |
+| `pack_256` | `inputs_embeds` [1, 256, 512], `attention_bias` [1, 1, 256, 256] (0 within a text, −1e4 elsewhere), `positions` [256, 1] (restart at 0 per text), `pool` [8, 256] (1/len over each text's tokens) | `embedding` [8, 768] |
+
+S ∈ {32, 48, 64, 128, 256, 512}. Inputs longer than 512 tokens must be truncated; on our test set capping at 512
+tokens changed nothing measurable. At ≤ 512 tokens every sliding-window layer sees the whole sequence, so a padding
+mask is all the model needs.
+
+`pack_256` runs up to eight short texts in one call. Short calls are bound by streaming the weights from memory, so
+one 256-token call with eight texts costs about the same as two single calls.
+
+## Use
+
+1. Prepend the task prefix (e.g. `title: none | text: ` for documents, `task: search result | query: ` for queries;
+   all in `config.json`).
+2. Tokenize with `tokenizer.json`: `<bos>` (2), tokens, `<eos>` (1).
+3. Look up each id's row in `embeddings.bf16` (bf16 → float: shift the 16 bits left by 16), multiply by √512 ≈
+   22.627417, pass as fp16.
+4. Call the smallest `embed_S` that fits, or pack several texts into `pack_256`.
+
+A Swift implementation (tokenizer, table lookup, packing, downloads) is in
+[FluidUse](https://github.com/FluidInference/FluidUse): `EmbeddingGemma2Manager`.
+
+Requires macOS 15 / iOS 18 (multifunction models). The first load on a device compiles all seven functions for the
+Neural Engine, which takes several minutes once; Core ML caches the result.
+
+## Numbers (M5 Pro, macOS 27)
+
+| | Result |
+|---|---|
+| Neural Engine placement | 3,862 / 3,862 ops (100%) in every function |
+| Latency, one text | 32 tokens 2.6 ms · 64 tokens 3.2 ms · 128 tokens 5.6 ms · 256 tokens 11.3 ms · 512 tokens 27.4 ms |
+| Throughput, short posts (`pack_256`, Swift) | 644 texts/s (one text per call: 307/s) |
+| vs PyTorch fp32 (sentence-transformers) | cosine ≥ 0.9984, mean 0.99996 over 562 posts |
+
+The model card warns against fp16: activations reach ~2,000, so squaring them inside RMSNorm overflows fp16. Every
+RMSNorm here divides by the row's max-abs value first, which keeps the math in range (plain fp16 RMSNorm: cosine
+0.95 to the reference; this conversion: 0.9996). The final normalization is rescaled the same way.
+
+License: Apache 2.0, same as the source model.

config.json

@@ -0,0 +1,27 @@
+{
+  "base_model": "google/embeddinggemma-2",
+  "format": "coreml",
+  "modality": "text",
+  "precision": "fp16",
+  "hidden_size": 512,
+  "embedding_dim": 768,
+  "matryoshka_dims": [768, 512, 256, 128],
+  "embed_scale": 22.627417,
+  "token_table": {"file": "embeddings.bf16", "dtype": "bfloat16", "shape": [262144, 512]},
+  "functions": {
+    "embed_32": {"tokens": 32}, "embed_48": {"tokens": 48}, "embed_64": {"tokens": 64}, "embed_128": {"tokens": 128},
+    "embed_256": {"tokens": 256}, "embed_512": {"tokens": 512},
+    "pack_256": {"tokens": 256, "slots": 8}
+  },
+  "bos_token_id": 2, "eos_token_id": 1, "pad_token_id": 0,
+  "prompts": {
+    "document": "title: none | text: ",
+    "search_query": "task: search result | query: ",
+    "question_answering": "task: question answering | query: ",
+    "fact_checking": "task: fact checking | query: ",
+    "code_retrieval": "task: code retrieval | query: ",
+    "classification": "task: classification | query: ",
+    "clustering": "task: clustering | query: ",
+    "sentence_similarity": "task: sentence similarity | query: "
+  }
+}

embeddings.bf16

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

tokenizer.json

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