python/smoke_test.py

#!/usr/bin/env python3
# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Smoke test: load the bundle, embed the shipped fictional texts, generated test images and generated audio clips,
compare with the fp32 reference vectors and check that every query finds its own image / clip first.

    python3 smoke_test.py [--bundle ../model] [--compute ane] [--image-compute gpu] [--audio-compute gpu] [--text-only]
                          [--no-audio] [--cleanup-temp] [--reference ../tests/smoke-reference.json ...]

Exit code 0 = every cosine >= its threshold (0.999 texts, 0.998 images, 0.9995 / 0.995 audio clips), all values
finite, every query at rank 1. `--text-only` runs only the v1.0 text reference and loads neither the vision nor the
audio asset; `--no-audio` skips the audio reference. The first run of a new executable (here: your python binary)
specializes the assets for the Neural Engine: expect a few minutes. Later runs load in about a second.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from embeddinggemma2_coreai import remove_own_mpsgraph_temp_files, run_smoke  # noqa: E402


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--bundle", default=str(HERE.parent / "model"))
    ap.add_argument("--reference", action="append", default=None,
                    help="reference file(s); default: tests/smoke-reference.json, image-smoke-reference.json and audio-smoke-reference.json")
    ap.add_argument("--compute", default="ane", choices=["ane", "gpu", "cpu", "default"])
    ap.add_argument("--image-compute", default=None, choices=["ane", "gpu", "cpu", "default"],
                    help="where the image encoder runs (default: same as --compute)")
    ap.add_argument("--audio-compute", default=None, choices=["ane", "gpu", "cpu", "default"],
                    help="where the audio encoder runs (default: same as --compute)")
    ap.add_argument("--text-only", action="store_true", help="only the text reference (the vision and audio assets are not loaded)")
    ap.add_argument("--no-audio", action="store_true", help="skip the audio reference (the audio asset is not loaded)")
    ap.add_argument("--cleanup-temp", action="store_true", help="remove this process's MPSGraph temp folders at exit")
    a = ap.parse_args()
    tests = HERE.parent / "tests"
    refs = a.reference or [str(tests / "smoke-reference.json")] + (
        [] if a.text_only else [str(tests / "image-smoke-reference.json")] + ([] if a.no_audio else [str(tests / "audio-smoke-reference.json")]))
    try:
        ok = run_smoke(a.bundle, refs, a.compute, a.image_compute, a.audio_compute)
    except ValueError as e:  # e.g. --audio-compute cpu (refused before anything is loaded): a message, not a traceback
        print(f"error: {e}", file=sys.stderr)
        return 2
    finally:
        if a.cleanup_temp:  # also after an error
            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")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())