tests/make_audio_smoke_reference.py

#!/usr/bin/env python3
# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Write tests/audio-smoke-reference.json: fp32 reference embeddings of the clips in tests/audio/ (the original model
through sentence-transformers, exactly as upstream does it: decode with libsndfile/ffmpeg, mono, libsoxr "HQ" to
16 kHz, no prompt) plus their captions as search queries and as documents.

    python3 make_audio_smoke_reference.py --model-dir <google/embeddinggemma-2 snapshot> [--audio audio] [--out audio-smoke-reference.json]

Needs torch, transformers >= 5.19, sentence-transformers >= 6.1, soundfile, soxr (and ffmpeg for AAC files).
Per-clip thresholds: clips already at 16 kHz and lossless keep the model-only tolerance (0.9995); clips the runners
resample with their own sinc filter instead of libsoxr, and MP3 files, get 0.995 (README "Audio: decoding and
resampling").
"""
from __future__ import annotations

import argparse
import json
import math
import shutil
import subprocess
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
Q_PROMPT = "task: search result | query: "
D_PROMPT = "title: none | text: "
# measured exceptions (README "Audio: accuracy"): digital silence is the least precise case of the Neural Engine's
# float16 arithmetic (0.9993 vs 0.9999 on the GPU); an MP3 with long near-silent decays amplifies the ~3e-5 difference
# between Apple's MP3 decoder (the runners) and libsndfile/mpg123 (this reference) in the log-mel floor (0.990 on both
# routes; the same clip decoded by libsndfile gives 0.9997).
THRESHOLD_OVERRIDES = {"edge_silence_1s_16k": (0.998, "digital silence: float16 arithmetic on the Neural Engine, ~0.9993"),
                       "long_melody_35s_16k": (0.985, "mp3 with silent decays: Apple's decoder vs libsndfile/mpg123, ~0.990 (0.9997 if decoded with libsndfile)")}


def decode(path: Path) -> tuple[np.ndarray, int]:
    import soundfile as sf
    try:
        x, sr = sf.read(str(path), dtype="float32", always_2d=True)
        return x, int(sr)
    except Exception:  # noqa: BLE001 - AAC: ffmpeg
        pass
    ffprobe = shutil.which("ffprobe")
    info = subprocess.run([ffprobe, "-v", "error", "-select_streams", "a:0", "-show_entries", "stream=sample_rate,channels",
                          "-of", "csv=p=0", str(path)], check=True, capture_output=True, text=True).stdout.strip().split(",")
    sr, ch = int(info[0]), int(info[1])
    raw = subprocess.run(["ffmpeg", "-nostdin", "-loglevel", "error", "-i", str(path), "-f", "f32le", "-acodec", "pcm_f32le", "-"],
                         check=True, capture_output=True).stdout
    return np.frombuffer(raw, dtype="<f4").reshape(-1, ch).copy(), sr


def prepare(path: Path) -> np.ndarray:
    import soxr
    x, sr = decode(path)
    mono = x[:, 0] if x.shape[1] == 1 else np.mean(x, axis=1, dtype=np.float32)
    if sr == 16000:
        return mono.astype(np.float32)
    y = soxr.resample(mono, sr, 16000, quality="soxr_hq").astype(np.float32)
    n_out = int(math.ceil(len(mono) * 16000 / sr))
    return np.pad(y, (0, max(0, n_out - len(y))))[:n_out]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", required=True)
    ap.add_argument("--audio", default=str(HERE / "audio"))
    ap.add_argument("--out", default=str(HERE / "audio-smoke-reference.json"))
    args = ap.parse_args()
    import torch
    from sentence_transformers import SentenceTransformer

    audio = Path(args.audio)
    manifest = json.loads((audio / "manifest.json").read_text(encoding="utf-8"))
    m = SentenceTransformer(args.model_dir, device="cpu", model_kwargs={"torch_dtype": torch.float32}, config_kwargs={"vision_config": None})
    clips, texts = [], []
    for it in manifest:
        wav = prepare(audio / it["file"])
        if len(wav) <= 128:
            continue  # the original model cannot embed fewer than 129 samples (the port gives the empty-sequence embedding)
        with torch.no_grad():
            e = m.encode([{"array": wav, "sampling_rate": 16000}], normalize_embeddings=True)[0]
        thr = 0.9995 if (it["sample_rate"] == 16000 and not it["file"].endswith(".mp3")) else 0.995
        note = None
        if it["name"] in THRESHOLD_OVERRIDES:
            thr, note = THRESHOLD_OVERRIDES[it["name"]]
        clips.append(dict(file=f"audio/{it['file']}", query=it["caption"], seconds=round(len(wav) / 16000, 3), threshold_cosine=thr,
                          note=note, embedding=[round(float(v), 6) for v in e]))
        print(f"{it['file']}: {len(wav) / 16000:.2f} s, threshold {thr}", flush=True)
    for it in manifest:
        if not it["caption"]:
            continue
        with torch.no_grad():
            q = m.encode([Q_PROMPT + it["caption"]], normalize_embeddings=True)[0]
            d = m.encode([D_PROMPT + it["caption"]], normalize_embeddings=True)[0]
        texts.append(dict(text=it["caption"], kind="search_query", embedding=[round(float(v), 6) for v in q]))
        texts.append(dict(text=it["caption"], kind="document", embedding=[round(float(v), 6) for v in d]))
    # which queries the fp32 reference itself ranks first (the runners check rank 1 only for those)
    cap = [(c, np.array(c["embedding"])) for c in clips if c["query"]]
    for c, _ in cap:
        q = next(np.array(t["embedding"]) for t in texts if t["kind"] == "search_query" and t["text"] == c["query"])
        scores = [(float(q @ a), cc["file"]) for cc, a in cap]
        c["fp32_rank1"] = max(scores)[1] == c["file"]
        print(f"fp32 ranks its own clip first: {c['fp32_rank1']}  {c['query']}", flush=True)
    out = dict(note="fp32 reference of google/embeddinggemma-2 through sentence-transformers (audio: libsndfile/ffmpeg + libsoxr HQ, no prompt); "
                    "per-clip thresholds: 0.9995 for 16 kHz lossless clips, 0.995 for resampled and MP3 clips, two measured exceptions "
                    "(note fields); fp32_rank1 = the fp32 model ranks the clip first for its own caption (synthesized sounds are a hard "
                    "retrieval case; the runners check rank 1 only where fp32 succeeds)",
               threshold_cosine=0.9995, threshold_cosine_audio=0.9995, audio=clips, texts=texts)
    Path(args.out).write_text(json.dumps(out, indent=1, ensure_ascii=False) + "\n", encoding="utf-8")
    print(f"{len(clips)} clips, {len(texts)} texts -> {args.out}")


if __name__ == "__main__":
    main()