#!/usr/bin/env python3 # Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE). """EmbeddingGemma 2 (text + images + audio) on Apple Core AI, macOS 27 — Python runner (library + CLI). Library: from embeddinggemma2_coreai import EmbeddingGemma2CoreAI, remove_own_mpsgraph_temp_files m = EmbeddingGemma2CoreAI("../model", compute="ane", load_vision=True, load_audio=True) # ane | gpu | cpu | default # images on the GPU (about 3x faster, more power), text on the Neural Engine: # m = EmbeddingGemma2CoreAI("../model", compute="ane", load_vision=True, image_compute="gpu") # audio on the GPU, independently: ..., load_audio=True, audio_compute="gpu" m.prewarm() # optional: one call per bucket (+ the vision/audio functions) q = m.embed("bar chart of lamp units sold per city", kind="search_query") # 768-d, L2-normalized v = m.embed_image("chart.png") # path or PIL image, same 768-d space a = m.embed_audio("clip.wav") # path (wav/aiff/flac/m4a/mp3/caf...) or (samples, rate) d = m.embed("The credits expire on the last day of March.", kind="document", dim=256) score = float(q @ v) del m remove_own_mpsgraph_temp_files() # see "Temp files" in the README `EmbeddingGemma2CoreAI(bundle, compute)` without `load_vision=True` / `load_audio=True` loads the text asset only (as in v1.0). `compute` places the text asset; `image_compute` / `audio_compute` place the image and audio encoders (default: the same as `compute`). CLI: python3 embeddinggemma2_coreai.py embed ../model [--compute ane] [--kind document] [--dim 768] "text" ... python3 embeddinggemma2_coreai.py embed-image ../model [--compute ane] [--image-compute gpu] [--exif] [--dim 768] image ... python3 embeddinggemma2_coreai.py embed-audio ../model [--compute ane] [--audio-compute gpu] [--long truncate|chunk_mean] [--dim 768] clip ... python3 embeddinggemma2_coreai.py smoke ../model ../tests/smoke-reference.json ../tests/image-smoke-reference.json ../tests/audio-smoke-reference.json [--image-compute gpu] [--audio-compute gpu] [--cleanup-temp] python3 embeddinggemma2_coreai.py tokcheck ../model ../tests/tokenizer-cases.json python3 embeddinggemma2_coreai.py latency ../model [--compute ane] [--image-compute gpu] [--audio-compute gpu] [--trials 50] [--text-only] python3 embeddinggemma2_coreai.py prewarm ../model [--compute ane] [--image-compute gpu] [--audio-compute gpu] [--text-only] Text goes through `mm_` with zero soft-token inputs (identical to the v1.0 text graph). Images are decoded with Pillow (`convert("RGB")`: alpha dropped, EXIF orientation not applied unless --exif), resized with the torchvision-exact bicubic in eg2_image.py and fed to `vision_2520`; the soft tokens then go through `mm_288` inside the sequence <|image> [n slots] . Audio is decoded, mixed to mono, resampled to 16 kHz and turned into log-mel frames by eg2_audio_frontend.py (numpy; soundfile or macOS afconvert for decoding), fed to `audio_` and then through `mm_` inside <|audio> [n slots] . Tokenization uses the Hugging Face `tokenizers` library on the bundle's tokenizer.json (the reference tokenizer). """ from __future__ import annotations import argparse import asyncio import ctypes import json import os import sys import time from pathlib import Path import numpy as np HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) import eg2_audio_frontend # noqa: E402 import eg2_image # noqa: E402 _PROCESS_START = time.time() class EmbeddingGemma2CoreAI: """Loads the bundle folder (config.json, tokenizer/, the text asset and optionally the vision and audio assets) and embeds texts, images and audio clips into the same 768-d space.""" def __init__(self, bundle: str | os.PathLike, compute: str = "ane", quiet: bool = True, load_vision: bool = False, image_compute: str | None = None, load_audio: bool = False, audio_compute: str | None = None): """compute: where the text asset runs (ane | gpu | cpu | default). image_compute / audio_compute: where the image and audio encoders run, independently (None = same as compute). "ane" uses the least CPU, GPU and power; "gpu" is about 3x faster per image (README, "Images: Neural Engine or GPU") and about 2x faster per 30 s of audio (README, "Audio: Neural Engine or GPU").""" import coreai.runtime as rt from tokenizers import Tokenizer self._rt = rt self.bundle = Path(bundle) self.config = json.loads((self.bundle / "config.json").read_text(encoding="utf-8")) self.hidden = self.config["hidden_size"] t = time.perf_counter() self.tokenizer = Tokenizer.from_file(str(self.bundle / "tokenizer" / "tokenizer.json")) self.tokenizer_load_s = time.perf_counter() - t self.compute = compute self.image_compute = image_compute or compute self.audio_compute = audio_compute or compute # refuse an unsupported audio request before loading anything (no specialization, no temp files) if load_audio and self.config.get("audio") is None: raise ValueError("this bundle has no audio asset (config.json 'audio')") if load_audio and self.audio_compute == "cpu": raise ValueError("audio_compute='cpu' is not supported for audio; use 'ane' or 'gpu'") opts = _specialization_options(rt, compute) image_opts = _specialization_options(rt, self.image_compute) audio_opts = _specialization_options(rt, self.audio_compute) t = time.perf_counter() self.model = asyncio.run(rt.AIModel.load(self.bundle / self.config["asset"], opts)) self.functions: dict[int, object] = {} for name in self.config["functions"]: self.functions[int(name.rsplit("_", 1)[1])] = self.model.load_function(name) self.buckets = sorted(self.functions) self.vision_model = self.vision_function = None if load_vision: vc = self.config["vision"] self.vision_model = asyncio.run(rt.AIModel.load(self.bundle / vc["asset"], image_opts)) self.vision_function = self.vision_model.load_function(vc["function"]) self.audio_model = None self.audio_functions: dict[int, object] = {} if load_audio: ac = self.config["audio"] self.audio_model = asyncio.run(rt.AIModel.load(self.bundle / ac["asset"], audio_opts)) for name in ac["functions"]: self.audio_functions[int(name.rsplit("_", 1)[1])] = self.audio_model.load_function(name) self.audio_buckets = sorted(self.audio_functions) self.load_s = time.perf_counter() - t if not quiet: print(f"loaded {self.bundle.name} compute={compute} image_compute={self.image_compute if load_vision else '-'} " f"audio_compute={self.audio_compute if load_audio else '-'} buckets={self.buckets} vision={load_vision} " f"audio={load_audio} in {self.load_s:.1f}s", flush=True) @property def has_vision(self) -> bool: return self.vision_function is not None @property def has_audio(self) -> bool: return bool(self.audio_functions) # ---- tokenization ------------------------------------------------------------------------------------- def tokenize(self, text: str, kind: str | None = "document") -> list[int]: """ + prompt prefix + text + , truncated to the largest bucket (longer input is cut).""" prefix = self.config["prompts"].get(kind, "") if kind else "" ids = self.tokenizer.encode(prefix + text, add_special_tokens=False).ids max_body = self.buckets[-1] - 2 if len(ids) > max_body: ids = ids[:max_body] return [self.config["bos_token_id"]] + ids + [self.config["eos_token_id"]] def bucket(self, n_tokens: int) -> int: for b in self.buckets: if b >= n_tokens: return b raise ValueError(f"{n_tokens} tokens exceeds the largest bucket {self.buckets[-1]}") # ---- text --------------------------------------------------------------------------------------------- def embed_ids(self, ids: list[int], bucket: int | None = None) -> np.ndarray: """Runs mm_ on ids (with /), right-padded with , no soft tokens. Returns the raw 768-d output (already L2-normalized by the graph).""" return self.embed_sequence(ids, bucket=bucket) def embed(self, text: str, kind: str | None = "document", dim: int = 768) -> np.ndarray: """L2-normalized embedding. dim 768, or a Matryoshka size (512/256/128): leading dims, re-normalized.""" return self._finish(self.embed_ids(self.tokenize(text, kind)), dim) # ---- images ------------------------------------------------------------------------------------------- def embed_image(self, image, apply_exif_orientation: bool = False, dim: int = 768) -> np.ndarray: """Path or PIL image -> L2-normalized 768-d embedding (or a Matryoshka prefix). The default apply_exif_orientation=False reproduces the upstream reference; True embeds camera photos upright.""" inputs = eg2_image.prepare_vision_inputs(image, apply_exif_orientation) ids, soft, slots = self.image_sequence(self.soft_tokens(inputs)) return self._finish(self.embed_sequence(ids, soft, slots), dim) def soft_tokens(self, inputs: eg2_image.VisionInputs) -> np.ndarray: """[n_soft, 512] float32 soft tokens of a prepared image (vision_2520).""" if self.vision_function is None: raise RuntimeError("vision asset not loaded: create the model with load_vision=True") rt = self._rt out = asyncio.run(self.vision_function(inputs={ "patches": rt.NDArray(inputs.patches), "pos_x": rt.NDArray(inputs.pos_x), "pos_y": rt.NDArray(inputs.pos_y), "cos": rt.NDArray(inputs.cos), "sin": rt.NDArray(inputs.sin), "patch_mask": rt.NDArray(inputs.patch_mask)})) return np.array(out["soft_tokens"].numpy(), dtype=np.float32)[0, : inputs.n_soft] def image_sequence(self, soft: np.ndarray) -> tuple[list[int], np.ndarray, list[int]]: """(ids, soft tokens, slots) of the upstream single-image sequence <|image> [n] .""" c = self.config n = soft.shape[0] ids = eg2_image.image_sequence(n, c["bos_token_id"], c["boi_token_id"], c["image_token_id"], c["eoi_token_id"], c["eos_token_id"]) return ids, soft, list(range(2, 2 + n)) # ---- audio -------------------------------------------------------------------------------------------- def embed_audio(self, audio, sample_rate: int | None = None, dim: int = 768, long_audio: str = "truncate") -> np.ndarray: """Path (anything soundfile or macOS afconvert decodes) or a numpy array of samples ([n] or [n, channels], with sample_rate) -> L2-normalized 768-d embedding (or a Matryoshka prefix). The clip is mixed to mono, resampled to 16 kHz and, with long_audio="truncate" (upstream behaviour), its first 30 s are embedded; "chunk_mean" averages the embeddings of consecutive 30 s windows instead. A channel-first array ([channels, n], the librosa/torchaudio layout) is rejected: pass audio.T.""" if isinstance(audio, (str, os.PathLike)): wav = eg2_audio_frontend.load_mono_16k(audio) else: if sample_rate is None or int(sample_rate) != sample_rate or sample_rate <= 0: raise ValueError(f"a positive integer sample_rate is required for sample arrays (got {sample_rate!r})") x = np.asarray(audio, dtype=np.float32) if x.ndim not in (1, 2): raise ValueError(f"expected samples of shape [n] or [n, channels], got {x.shape}") if x.ndim == 2 and x.shape[1] > max(x.shape[0], 1): raise ValueError(f"samples of shape {x.shape} look channel-first ([channels, n], the librosa/torchaudio " "layout): pass [n, channels], e.g. audio.T") wav = eg2_audio_frontend.resample_16k(eg2_audio_frontend.to_mono(x), int(sample_rate)) return self.embed_audio_16k(wav, dim, long_audio) def embed_audio_16k(self, wav16k: np.ndarray, dim: int = 768, long_audio: str = "truncate") -> np.ndarray: """16 kHz mono float32 samples -> embedding.""" if long_audio == "truncate": return self._finish(self.embed_audio_inputs(eg2_audio_frontend.prepare_audio_inputs(wav16k, self.audio_buckets)), dim) if long_audio != "chunk_mean": raise ValueError("long_audio must be truncate or chunk_mean") acc = np.zeros(self.config["embedding_dim"], np.float64) for chunk in eg2_audio_frontend.chunks_30s(wav16k): acc += self.embed_audio_inputs(eg2_audio_frontend.prepare_audio_inputs(chunk, self.audio_buckets)) return self._finish(acc.astype(np.float32), dim) def embed_audio_inputs(self, inputs: eg2_audio_frontend.AudioInputs) -> np.ndarray: """Prepared inputs -> raw 768-d output of mm_ (already L2-normalized by the graph).""" soft = self.audio_soft_tokens(inputs) ids, soft, slots = self.audio_sequence(soft) return self.embed_sequence(ids, soft if slots else None, slots or None) def audio_soft_tokens(self, inputs: eg2_audio_frontend.AudioInputs) -> np.ndarray: """[n_tokens, 512] float32 soft tokens of prepared audio (audio_).""" fn = self.audio_functions.get(inputs.bucket) if fn is None: raise RuntimeError("audio asset not loaded: create the model with load_audio=True" if not self.audio_functions else f"no audio function for bucket {inputs.bucket}") rt = self._rt out = asyncio.run(fn(inputs={"features": rt.NDArray(inputs.features), "frame_mask": rt.NDArray(inputs.frame_mask), "half_mask": rt.NDArray(inputs.half_mask), "token_mask": rt.NDArray(inputs.token_mask)})) return np.array(out["soft_tokens"].numpy(), dtype=np.float32)[0, : inputs.n_tokens] def audio_sequence(self, soft: np.ndarray) -> tuple[list[int], np.ndarray, list[int]]: """(ids, soft tokens, slots) of the upstream single-clip sequence <|audio> [n] (no prompt).""" c = self.config n = soft.shape[0] ids = eg2_audio_frontend.audio_sequence(n, c["bos_token_id"], c.get("boa_token_id", 256000), c.get("audio_token_id", 258881), c.get("eoa_token_id", 258883), c["eos_token_id"]) return ids, soft, list(range(2, 2 + n)) # ---- any sequence ------------------------------------------------------------------------------------- def embed_sequence(self, ids: list[int], soft: np.ndarray | None = None, slots: list[int] | None = None, bucket: int | None = None) -> np.ndarray: """Runs mm_ on ids (with /), soft tokens [len(slots), 512] placed at `slots` (any mix of text and images that fits a bucket). Returns the raw 768-d output (already L2-normalized by the graph).""" S = bucket or self.bucket(len(ids)) if len(ids) > S: raise ValueError(f"{len(ids)} ids do not fit bucket {S}") rt = self._rt x = np.full((1, S), self.config["pad_token_id"], dtype=np.int32) x[0, : len(ids)] = ids am = np.zeros((1, S), dtype=np.float16) am[0, : len(ids)] = 1 mm = np.zeros((1, S, self.hidden), dtype=np.float16) mmm = np.zeros((1, S), dtype=np.float16) if slots: if soft is None or soft.shape != (len(slots), self.hidden): raise ValueError("soft tokens do not match slots") mm[0, slots] = soft.astype(np.float16) mmm[0, slots] = 1 out = asyncio.run(self.functions[S](inputs={"input_ids": rt.NDArray(x), "attention_mask": rt.NDArray(am), "mm_embeds": rt.NDArray(mm), "mm_mask": rt.NDArray(mmm)})) return np.array(out["embedding"].numpy(), dtype=np.float32).reshape(-1) def _finish(self, v: np.ndarray, dim: int) -> np.ndarray: if dim not in self.config["matryoshka_dims"]: raise ValueError(f"dim must be one of {self.config['matryoshka_dims']}") v = v[:dim] return v / max(float(np.linalg.norm(v)), 1e-12) def prewarm(self) -> float: """One call per bucket (and one vision call, if loaded) so the first real call pays no first-call cost.""" t = time.perf_counter() for S in self.buckets: self.embed_ids([self.config["bos_token_id"], self.config["eos_token_id"]], S) if self.vision_function is not None: self.soft_tokens(eg2_image.patchify_cell_major(np.full((48, 48, 3), 128, np.uint8))) for tb in self.audio_buckets: self.audio_soft_tokens(eg2_audio_frontend.prepare_audio_inputs(np.zeros(min(tb, 3072) * 160, np.float32), [tb])) return time.perf_counter() - t def close(self) -> None: self.functions.clear() self.audio_functions.clear() self.model = self.vision_model = self.vision_function = self.audio_model = None def _specialization_options(rt, compute: str): if compute == "cpu": return rt.SpecializationOptions.cpu_only() if compute == "gpu": return rt.SpecializationOptions.from_preferred_compute_unit_kind(rt.ComputeUnitKind.gpu()) if compute == "ane": return rt.SpecializationOptions.from_preferred_compute_unit_kind(rt.ComputeUnitKind.neural_engine()) if compute == "default": return rt.SpecializationOptions.default() raise ValueError("compute must be ane | gpu | cpu | default") # ---- temp-file hygiene ---------------------------------------------------------------------------------- def mpsgraph_temp_root() -> Path: import tempfile return Path(os.environ.get("TMPDIR") or tempfile.gettempdir()) / "com.apple.MetalPerformanceShadersGraph" def remove_own_mpsgraph_temp_files(pid: int | None = None, since: float | None = None) -> tuple[int, int]: """Each Neural Engine load makes MPSGraph write compiled-graph copies into $TMPDIR/com.apple.MetalPerformance ShadersGraph/mpsgraph---/ (~360 MB for the text asset, ~300 MB for the vision asset), and the Python runtime (coreai.runtime, macOS 27) does NOT delete them when the model is released or the interpreter exits. This removes the folders of the CURRENT process only (its pid, created after it started). Call it after the model object has been released, or at exit. Returns (folders, bytes).""" import shutil pid = pid or os.getpid() since = _PROCESS_START - 2 if since is None else since root = mpsgraph_temp_root() folders = nbytes = 0 if not root.is_dir(): return 0, 0 for d in root.iterdir(): if not d.name.startswith(f"mpsgraph-{pid}-") or not d.is_dir(): continue try: if d.stat().st_birthtime < since: continue # left by an older process that happened to have the same pid nbytes += sum(f.stat().st_size for f in d.rglob("*") if f.is_file()) shutil.rmtree(d) folders += 1 except OSError: pass return folders, nbytes # ---- process accounting helpers (macOS) ------------------------------------------------------------------ def _rusage_v6(): class RU6(ctypes.Structure): _fields_ = [("uuid", ctypes.c_uint8 * 16)] + [(f"f{i}", ctypes.c_uint64) for i in range(47)] + [("reserved", ctypes.c_uint64 * 9)] libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib") r = RU6() return r if libc.proc_pid_rusage(os.getpid(), 6, ctypes.byref(r)) == 0 else None def footprints_mb() -> tuple[float, float]: """(lifetime max physical footprint, Neural Engine footprint) in MB; the latter is > 0 only if the ANE ran.""" r = _rusage_v6() if r is None: return -1.0, -1.0 return r.f28 / 2**20, r.f44 / 2**20 # ri_lifetime_max_phys_footprint, ri_neural_footprint def cosine(a, b) -> float: a = np.asarray(a, np.float64) b = np.asarray(b, np.float64) return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b))) # ---- smoke test (used by the CLI and smoke_test.py) ------------------------------------------------------- def run_smoke(bundle, ref_paths: list, compute: str = "ane", image_compute: str | None = None, audio_compute: str | None = None) -> bool: """Runs every reference file: text files (v1.0 format, "cases"), image files ("images" + "texts") and audio files ("audio" + "texts"). Returns True when every cosine is above its threshold, all values are finite and every query ranks its own image / clip first.""" refs = [(Path(p), json.loads(Path(p).read_text(encoding="utf-8"))) for p in ref_paths] need_vision = any(r.get("images") for _, r in refs) need_audio = any(r.get("audio") for _, r in refs) m = EmbeddingGemma2CoreAI(bundle, compute, load_vision=need_vision, image_compute=image_compute, load_audio=need_audio, audio_compute=audio_compute) print(f"loaded in {m.load_s:.1f} s (tokenizer {m.tokenizer_load_s:.2f} s), compute={compute}, " f"image_compute={m.image_compute if need_vision else '-'}, audio_compute={m.audio_compute if need_audio else '-'}, " f"buckets={m.buckets}, vision {'loaded' if m.has_vision else 'not loaded'}, audio {'loaded' if m.has_audio else 'not loaded'}") ok = True for path, ref in refs: print(f"-- {path.name}") thr = ref["threshold_cosine"] for c in ref.get("cases", []): v = m.embed(c["text"], c["kind"]) cos = cosine(v, c["embedding"]) passed = cos >= thr and bool(np.all(np.isfinite(v))) ok = ok and passed print(f"{'PASS' if passed else 'FAIL'} cosine {cos:.6f} (threshold {thr}) {c['kind']}: {c['text'][:60]}") if len(ref.get("cases", [])) > 1: c = ref["cases"][1] # Matryoshka prefixes vs the truncated + re-normalized reference for dim in (512, 256, 128): cos = cosine(m.embed(c["text"], c["kind"], dim=dim), c["embedding"][:dim]) ok = ok and cos >= thr print(f"{'PASS' if cos >= thr else 'FAIL'} cosine {cos:.6f} at dim {dim}") images = ref.get("images", []) thr_img = ref.get("threshold_cosine_image", thr) img_vecs = [] for c in images: t0 = time.perf_counter() v = m.embed_image(path.parent / c["file"]) cos = cosine(v, c["embedding"]) passed = cos >= thr_img and bool(np.all(np.isfinite(v))) ok = ok and passed img_vecs.append(v) print(f"{'PASS' if passed else 'FAIL'} cosine {cos:.6f} (threshold {thr_img}) image {c['file']} {int((time.perf_counter() - t0) * 1000)} ms") audio = ref.get("audio", []) aud_vecs = [] for c in audio: t0 = time.perf_counter() v = m.embed_audio(path.parent / c["file"]) cos = cosine(v, c["embedding"]) thr_a = c.get("threshold_cosine", ref.get("threshold_cosine_audio", thr)) passed = cos >= thr_a and bool(np.all(np.isfinite(v))) ok = ok and passed aud_vecs.append(v) print(f"{'PASS' if passed else 'FAIL'} cosine {cos:.6f} (threshold {thr_a}) audio {c['file']} {int((time.perf_counter() - t0) * 1000)} ms") queries = [] for c in ref.get("texts", []): v = m.embed(c["text"], c["kind"]) cos = cosine(v, c["embedding"]) passed = cos >= thr and bool(np.all(np.isfinite(v))) ok = ok and passed if c["kind"] == "search_query": queries.append((c["text"], v)) print(f"{'PASS' if passed else 'FAIL'} cosine {cos:.6f} (threshold {thr}) {c['kind']}: {c['text'][:60]}") if images: hits = 0 for q, qv in queries: scores = [cosine(qv, iv) for iv in img_vecs] best = int(np.argmax(scores)) want = next((i for i, c in enumerate(images) if c["query"] == q), -1) hits += best == want print(f"{'PASS' if best == want else 'FAIL'} query \"{q}\" -> {images[best]['file']} ({scores[best]:.3f})") ok = ok and hits == len(queries) print(f"text -> image: {hits}/{len(queries)} queries rank their own image first") if audio: with_q = [i for i, c in enumerate(audio) if c.get("query")] hits = total = 0 for q, qv in queries: want = next((i for i, c in enumerate(audio) if c.get("query") == q), -1) if want < 0: continue scores = [cosine(qv, aud_vecs[i]) for i in with_q] best = with_q[int(np.argmax(scores))] if audio[want].get("fp32_rank1") is False: # checked only where fp32 itself ranks the clip first print(f"SKIP query \"{q}\" -> {audio[best]['file']} ({max(scores):.3f}) (fp32 does not rank this clip first either)") continue total += 1 hits += best == want print(f"{'PASS' if best == want else 'FAIL'} query \"{q}\" -> {audio[best]['file']} ({max(scores):.3f})") ok = ok and hits == total print(f"text -> audio: {hits}/{total} queries rank their own clip first (as fp32 does)") fp, nf = footprints_mb() print(f"peak footprint {fp:.0f} MB, neural footprint {nf:.0f} MB ({'Neural Engine in use' if nf > 0 else 'Neural Engine not used'})") print("SMOKE PASS" if ok else "SMOKE FAIL") m.close() del m return ok # ---- CLI ------------------------------------------------------------------------------------------------- def cmd_embed(a) -> int: m = EmbeddingGemma2CoreAI(a.bundle, a.compute) for t in a.args: v = m.embed(t, None if a.kind == "none" else a.kind, a.dim) print(json.dumps({"text": t, "kind": a.kind, "dim": int(v.shape[0]), "embedding": [round(float(x), 6) for x in v]}, ensure_ascii=False)) m.close() return 0 def cmd_embed_image(a) -> int: m = EmbeddingGemma2CoreAI(a.bundle, a.compute, load_vision=True, image_compute=a.image_compute) for p in a.args: t0 = time.perf_counter() inputs = eg2_image.prepare_vision_inputs(p, a.exif) t1 = time.perf_counter() ids, soft, slots = m.image_sequence(m.soft_tokens(inputs)) v = m._finish(m.embed_sequence(ids, soft, slots), a.dim) print(json.dumps({"image": p, "dim": a.dim, "n_soft": inputs.n_soft, "resized": [inputs.resized[1], inputs.resized[0]], "exif_orientation_applied": a.exif, "preprocess_ms": int((t1 - t0) * 1000), "model_ms": int((time.perf_counter() - t1) * 1000), "embedding": [round(float(x), 6) for x in v]})) m.close() return 0 def cmd_embed_audio(a) -> int: m = EmbeddingGemma2CoreAI(a.bundle, a.compute, load_audio=True, audio_compute=a.audio_compute) for p in a.args: t0 = time.perf_counter() d = eg2_audio_frontend.decode(p) wav = eg2_audio_frontend.resample_16k(eg2_audio_frontend.to_mono(d.samples), d.sample_rate) t1 = time.perf_counter() v = m.embed_audio_16k(wav, a.dim, a.long) frames = eg2_audio_frontend.num_valid_frames(len(wav)) print(json.dumps({"audio": p, "dim": a.dim, "sample_rate": d.sample_rate, "channels": int(d.samples.shape[1]), "decoder": d.decoder, "seconds": round(len(wav) / 16000, 3), "n_frames": frames, "n_soft": eg2_audio_frontend.n_audio_tokens(frames), "truncated_to_30s": bool(len(wav) > eg2_audio_frontend.MAX_SAMPLES and a.long == "truncate"), "long_audio": a.long, "decode_ms": int((t1 - t0) * 1000), "model_ms": int((time.perf_counter() - t1) * 1000), "embedding": [round(float(x), 6) for x in v]})) m.close() return 0 def cmd_smoke(a) -> int: if not a.args: raise SystemExit("smoke needs at least one reference file") return 0 if run_smoke(a.bundle, a.args, a.compute, a.image_compute, a.audio_compute) else 1 def cmd_tokcheck(a) -> int: from tokenizers import Tokenizer cases = json.loads(Path(a.args[0]).read_text(encoding="utf-8")) tok = Tokenizer.from_file(str(Path(a.bundle) / "tokenizer" / "tokenizer.json")) prefixes = {"q": "task: search result | query: ", "d": "title: none | text: ", "raw": ""} same = 0 t = time.perf_counter() for c in cases: ids = [2] + tok.encode(prefixes[c["kind"]] + c["text"], add_special_tokens=False).ids + [1] same += ids == c["ids"] print(f"TOKCHECK {same}/{len(cases)} identical (tokenizers {__import__('tokenizers').__version__}) in {time.perf_counter() - t:.2f} s") return 0 if same == len(cases) else 1 def _has_audio_asset(bundle) -> bool: return "audio" in json.loads((Path(bundle) / "config.json").read_text(encoding="utf-8")) def cmd_latency(a) -> int: m = EmbeddingGemma2CoreAI(a.bundle, a.compute, load_vision=not a.text_only, image_compute=a.image_compute, load_audio=not a.text_only and _has_audio_asset(a.bundle), audio_compute=a.audio_compute) print(f"loaded in {m.load_s:.1f} s, compute={a.compute}, image_compute={m.image_compute if m.has_vision else '-'}, " f"audio_compute={m.audio_compute if m.has_audio else '-'}, vision {'loaded' if m.has_vision else 'not loaded'}, " f"audio {'loaded' if m.has_audio else 'not loaded'}; prewarm {m.prewarm():.2f} s") if m.has_audio: n_a = max(5, a.trials // 2) print("audio seconds bucket tokens audio_ p50 ms mm_ p50 ms end-to-end p50 ms CPU ms/clip front end ms") for seconds in (3.0, 10.0, 30.0): t = np.arange(int(seconds * 16000)) / 16000.0 wav = (0.3 * np.sin(2 * np.pi * 220 * t) + 0.2 * np.sin(2 * np.pi * 1375 * t + 0.5 * np.sin(2 * np.pi * 3 * t)) + 0.05 * ((np.arange(len(t)) * 7919) % 1000) / 1000 - 0.025).astype(np.float32) tf0 = time.perf_counter() inputs = eg2_audio_frontend.prepare_audio_inputs(wav, m.audio_buckets) front_ms = (time.perf_counter() - tf0) * 1000 for _ in range(3): m.embed_audio_inputs(inputs) ta, tt, te = [], [], [] c0 = time.process_time() for _ in range(n_a): t0 = time.perf_counter() soft = m.audio_soft_tokens(inputs) t1 = time.perf_counter() ids, soft, slots = m.audio_sequence(soft) m.embed_sequence(ids, soft, slots) t2 = time.perf_counter() ta.append((t1 - t0) * 1000); tt.append((t2 - t1) * 1000); te.append((t2 - t0) * 1000) cpu = (time.process_time() - c0) / n_a * 1000 print(f" {seconds:7.1f} {inputs.bucket:6d} {inputs.n_tokens:6d} {np.median(ta):17.1f} {np.median(tt):14.1f} " f"{np.median(te):18.1f} {cpu:12.1f} {front_ms:12.1f}") n48 = 30 * 48000 st = np.stack([np.sin(np.arange(n48) * 0.01), np.sin(np.arange(n48) * 0.013)], 1).astype(np.float32) tp0 = time.perf_counter() w16 = eg2_audio_frontend.resample_16k(eg2_audio_frontend.to_mono(st), 48000) tp1 = time.perf_counter() eg2_audio_frontend.prepare_audio_inputs(w16, m.audio_buckets) print(f"audio front end, 30 s stereo 48 kHz (numpy): mono + resample {(tp1 - tp0) * 1000:.0f} ms, log-mel + inputs {(time.perf_counter() - tp1) * 1000:.0f} ms") if m.has_vision: rng = np.random.RandomState(0) px = rng.randint(0, 256, (624, 1008, 3)).astype(np.uint8) inputs = eg2_image.patchify_cell_major(px) tv, tt = [], [] n = max(5, a.trials // 2) c0 = time.process_time() for _ in range(n): t0 = time.perf_counter() soft = m.soft_tokens(inputs) t1 = time.perf_counter() ids, soft, slots = m.image_sequence(soft) m.embed_sequence(ids, soft, slots) tv.append((t1 - t0) * 1000) tt.append((time.perf_counter() - t1) * 1000) cpu = (time.process_time() - c0) / n * 1000 print(f"image (273 soft tokens, {n} trials): vision_2520 p50 {np.median(tv):.1f} ms (min {min(tv):.1f}), mm_288 p50 {np.median(tt):.1f} ms, " f"end-to-end p50 {np.median(np.array(tv) + np.array(tt)):.1f} ms, CPU {cpu:.1f} ms/image (excl. decode+resize)") big = rng.randint(0, 256, (1600, 2560, 3)).astype(np.uint8) t0 = time.perf_counter() for _ in range(3): eg2_image.patchify_cell_major(eg2_image.resize_bicubic_aa_u8(big, 624, 1008)) print(f"preprocess (numpy resize 2560x1600 -> 1008x624 + patches): {(time.perf_counter() - t0) / 3 * 1000:.0f} ms") print("bucket p50 ms min ms p90 ms CPU ms/call tok/s") for S in m.buckets: ids = [m.config["bos_token_id"]] + [1000 + (i * 7919) % 200000 for i in range(1, S - 1)] + [m.config["eos_token_id"]] for _ in range(3): m.embed_ids(ids, S) ts = [] c0 = time.process_time() for _ in range(a.trials): t = time.perf_counter() m.embed_ids(ids, S) ts.append((time.perf_counter() - t) * 1000) cpu = (time.process_time() - c0) / a.trials * 1000 p50 = float(np.median(ts)) print(f"{S:5d} {p50:7.2f} {min(ts):7.2f} {float(np.percentile(ts, 90)):7.2f} {cpu:11.2f} {S / p50 * 1000:6.0f}") fp, nf = footprints_mb() print(f"peak footprint {fp:.0f} MB, neural footprint {nf:.0f} MB ({'Neural Engine in use' if nf > 0 else 'Neural Engine NOT used'})") m.close() return 0 def cmd_prewarm(a) -> int: t0 = time.perf_counter() m = EmbeddingGemma2CoreAI(a.bundle, a.compute, load_vision=not a.text_only, image_compute=a.image_compute, load_audio=not a.text_only and _has_audio_asset(a.bundle), audio_compute=a.audio_compute) w = m.prewarm() print(f"load {m.load_s:.1f} s + prewarm {w:.2f} s = ready in {time.perf_counter() - t0:.1f} s " f"(vision {'loaded' if m.has_vision else 'not loaded'}, audio {'loaded' if m.has_audio else 'not loaded'})") m.close() return 0 def main() -> None: ap = argparse.ArgumentParser(description="EmbeddingGemma 2 (text + images + audio) on Apple Core AI") ap.add_argument("cmd", choices=["embed", "embed-image", "embed-audio", "smoke", "tokcheck", "latency", "prewarm"]) ap.add_argument("bundle", help="the model folder (config.json, tokenizer/, *.aimodel)") ap.add_argument("args", nargs="*", help="texts (embed), image files (embed-image), audio files (embed-audio) or json files (smoke, tokcheck)") ap.add_argument("--compute", default="ane", choices=["ane", "gpu", "cpu", "default"], help="text asset (and the encoders, unless --image-compute / --audio-compute)") ap.add_argument("--image-compute", default=None, choices=["ane", "gpu", "cpu", "default"], help="image encoder only (default: same as --compute); gpu is ~3x faster per image, ane uses the least power") ap.add_argument("--audio-compute", default=None, choices=["ane", "gpu", "cpu", "default"], help="audio encoder only (default: same as --compute)") ap.add_argument("--kind", default="document", help="prompt kind from config.json (document, search_query, ...) or none") ap.add_argument("--exif", action="store_true", help="embed-image: apply the EXIF orientation first (upstream does not)") ap.add_argument("--long", default="truncate", choices=["truncate", "chunk_mean"], help="embed-audio: audio longer than 30 s") ap.add_argument("--dim", type=int, default=768) ap.add_argument("--trials", type=int, default=50) ap.add_argument("--text-only", action="store_true", help="latency, prewarm: do not load the vision and audio assets") ap.add_argument("--cleanup-temp", action="store_true", help="remove this process's MPSGraph temp folders at exit") a = ap.parse_args() try: code = {"embed": cmd_embed, "embed-image": cmd_embed_image, "embed-audio": cmd_embed_audio, "smoke": cmd_smoke, "tokcheck": cmd_tokcheck, "latency": cmd_latency, "prewarm": cmd_prewarm}[a.cmd](a) except ValueError as e: # bad input or an unsupported route: a message, not a traceback print(f"error: {e}", file=sys.stderr) code = 2 finally: if a.cleanup_temp: # also after an error, so that a failed run does not leave its temp folders behind import gc gc.collect() n, b = remove_own_mpsgraph_temp_files() print(f"cleanup-temp: removed {n} MPSGraph temp folder(s), {b / 2**20:.1f} MB") sys.exit(code) if __name__ == "__main__": main()