Add audio: EmbeddingGemma2Audio (USM conformer, log-mel in-graph, GPU fp16, 617x real time)

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

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

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

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

EmbeddingGemma2Audio.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "33C920C7-8DF5-410F-BC71-B016FF1BDAB6": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "53956E94-BB15-4BE5-A123-E22EABB7D9A9": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "53956E94-BB15-4BE5-A123-E22EABB7D9A9"
+}

README.md

@@ -24,6 +24,9 @@ Output: the same 768-d, L2-normalized embedding as `SentenceTransformer("google/
 (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.
 
+**Audio** (added): `EmbeddingGemma2Audio.mlpackage` turns 10 s of 16 kHz audio into 250 tokens in the text model's
+space; the text functions then embed them, so audio and text share one space (search audio with a text query).
+
 ## Files
 
 | File | What |
@@ -31,7 +34,8 @@ the first N values and re-normalize.
 | `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 |
+| `EmbeddingGemma2Audio.mlpackage` | audio encoder (Gemma 4 USM conformer, 12 layers) + projection, fp16, log-mel in-graph (561 MB) |
+| `config.json` | shapes, function list, token ids, task prefixes, audio window layout |
 
 ## Functions
 
@@ -50,6 +54,25 @@ 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.
 
+## Audio
+
+`EmbeddingGemma2Audio.mlpackage`: `waveform` [1, 160160] fp32 (160 zero samples, then up to 10 s of 16 kHz mono,
+zero-padded) and `frame_mask` [1, 1000] (frame i valid when `i*160 + 321 <= 160 + samples`) → `audio_tokens`
+[1, 250, 512]; token k is valid when `frame_mask[4k]` is. Build `<bos> <|audio> tokens <audio|> <eos>` (special
+tokens from `embeddings.bf16` × √512, audio tokens as they are) and run `embed_256`. Log-mel (Gemma 4 feature
+extractor: 20 ms Hann frames, 10 ms hop, 128 HTK mel bins, log(mag + 1e-3)) is computed inside the model.
+
+Run it on the **GPU** (`cpuAndGPU`): the chunked local attention's 5-D blocks fall back to the CPU on the Neural
+Engine (46 ms per window there vs 13.6 ms on the GPU). fp16 drifts from fp32 on a few quiet tokens (window embedding
+cos ≥ 0.977, mean 0.9993); on a 1-hour earnings-call index, the top hit matched fp32 for 15/15 text queries and the
+top-5 overlap was 95%.
+
+| 1 hour of audio, M5 Pro | Time |
+|---|---|
+| audio model (GPU, fp16) | 5.8 s (617× real time) |
+| audio model + text model, one window after another | 10.2 s (353×) |
+| fp32 audio model (exact to PyTorch, token cos 1.0000) | 13.2 s (272×) |
+
 ## Use
 
 1. Prepend the task prefix (e.g. `title: none | text: ` for documents, `task: search result | query: ` for queries;

config.json

@@ -1,19 +1,55 @@
 {
   "base_model": "google/embeddinggemma-2",
   "format": "coreml",
-  "modality": "text",
+  "modality": [
+    "text",
+    "audio"
+  ],
   "precision": "fp16",
   "hidden_size": 512,
   "embedding_dim": 768,
-  "matryoshka_dims": [768, 512, 256, 128],
+  "matryoshka_dims": [
+    768,
+    512,
+    256,
+    128
+  ],
   "embed_scale": 22.627417,
-  "token_table": {"file": "embeddings.bf16", "dtype": "bfloat16", "shape": [262144, 512]},
+  "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}
+    "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,
+  "bos_token_id": 2,
+  "eos_token_id": 1,
+  "pad_token_id": 0,
   "prompts": {
     "document": "title: none | text: ",
     "search_query": "task: search result | query: ",
@@ -23,5 +59,21 @@
     "classification": "task: classification | query: ",
     "clustering": "task: clustering | query: ",
     "sentence_similarity": "task: sentence similarity | query: "
+  },
+  "audio": {
+    "package": "EmbeddingGemma2Audio.mlpackage",
+    "compute_units": "cpuAndGPU",
+    "precision": "fp16",
+    "sample_rate": 16000,
+    "window_seconds": 10,
+    "waveform_samples": 160160,
+    "left_pad_samples": 160,
+    "frames": 1000,
+    "frame_valid_if": "i * 160 + 321 <= 160 + audio_samples",
+    "tokens_per_window": 250,
+    "token_ms": 40,
+    "boa_token_id": 256000,
+    "eoa_token_id": 258883,
+    "text_sequence": "<bos> <|audio> [audio tokens] <audio|> <eos> -> embed_256"
   }
-}
+}