python/embeddinggemma2_coreai.py

#!/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_<S>` 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 <bos> <|image> [n slots] <image|> <eos>. 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_<Tb>` and then through `mm_<S>` inside <bos> <|audio> [n slots] <audio|> <eos>. 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]:
        """<bos> + prompt prefix + text + <eos>, 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_<bucket> on ids (with <bos>/<eos>), right-padded with <pad>, 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 <bos> <|image> [n] <image|> <eos>."""
        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_<S> (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_<bucket>)."""
        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 <bos> <|audio> [n] <audio|> <eos> (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_<bucket> on ids (with <bos>/<eos>), 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-<pid>-<date>-<n>/ (~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_<Tb> p50 ms   mm_<S> 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()