jaisawhney/SonicEmbed-ONNX

SonicEmbed ONNX

ONNX builds of SonicEmbed for ONNX Runtime and Transformers.js. Its model card covers the full results and training.

Build	Files	Size	R@10, three clips	R@10, one clip	R@1, three clips	Median rank, three clips
fp32	no suffix	2,257 MB	40.7	36.7	9.4	16.5
fp16	_fp16	1,128 MB	40.7	36.8	9.5	16.5
q8	_quantized	654 MB	40.6	36.6	9.6	16
q8 audio, q4 text	_quantized, _q4	514 MB	40.4	35.9	9.3	17
q4	_q4	363 MB	37.8	34.5	7.8	19
PyTorch	SonicEmbed	2,307 MB	40.4	36.7	9.5	16

Scores are text-to-music search on Song Describer. Gaps of a few tenths are noise, not real differences from PyTorch.

• Audio must already be 16 kHz mono: resample first.
• Only the first 30 s of audio is used: anything longer is truncated. To cover a whole track, average clips, as shown on SonicEmbed's card.
• Prompts: start text you search with task: search result | query:, and text you search through with title: none | text:. Audio takes no prompt.
• Not tested: general text quality (MTEB) and GPU backends weren't measured on these builds.

Usage

Transformers.js

import { AutoModel, AutoProcessor } from "@huggingface/transformers";

const id = "jaisawhney/SonicEmbed-ONNX";
const processor = await AutoProcessor.from_pretrained(id);
const model = await AutoModel.from_pretrained(id, { dtype: "fp16" });

const queryInputs = await processor("task: search result | query: calm piano");
const queryOutputs = await model(queryInputs);
const query = queryOutputs.sentence_embedding;

// audio is a mono 16 kHz Float32Array.
const trackInputs = await processor(null, null, audio);
const trackOutputs = await model(trackInputs);
const track = trackOutputs.sentence_embedding;

ONNX Runtime

pip install onnxruntime transformers huggingface_hub librosa

import librosa
import numpy as np
import onnxruntime as ort
from huggingface_hub import snapshot_download
from transformers import AutoFeatureExtractor, AutoTokenizer

BUILD = "_fp16"
NO_FEATURES = np.zeros((0, 512), "float32")

folder = snapshot_download(
    "jaisawhney/SonicEmbed-ONNX",
    allow_patterns=[f"onnx/audio_encoder{BUILD}.onnx*", f"onnx/model{BUILD}.onnx*", "*.json"],
    local_dir="SonicEmbed-ONNX",
)

tokenizer = AutoTokenizer.from_pretrained(folder)
extractor = AutoFeatureExtractor.from_pretrained(folder)
audio_encoder = ort.InferenceSession(f"{folder}/onnx/audio_encoder{BUILD}.onnx")
text_model = ort.InferenceSession(f"{folder}/onnx/model{BUILD}.onnx")


def embed(input_ids, audio_features):
    outputs = text_model.run(
        ["sentence_embedding"],
        {
            "input_ids": input_ids,
            "attention_mask": np.ones_like(input_ids),
            "image_features": NO_FEATURES,
            "video_features": NO_FEATURES,
            "audio_features": audio_features,
        },
    )
    embedding = outputs[0][0]
    return embedding / np.linalg.norm(embedding)


def embed_text(text):
    input_ids = tokenizer(text, return_tensors="np")["input_ids"]
    return embed(input_ids, NO_FEATURES)


def embed_audio(wave):
    inputs = extractor([wave], sampling_rate=16000, return_tensors="np")
    outputs = audio_encoder.run(
        None,
        {
            "input_features": inputs["input_features"],
            "input_features_mask": inputs["input_features_mask"],
        },
    )
    audio_features = outputs[0]

    audio_start, audio_token, audio_end = tokenizer.convert_tokens_to_ids(
        ["<|audio>", "<|audio|>", "<audio|>"]
    )
    tokens = [tokenizer.bos_token_id, audio_start]
    tokens += [audio_token] * len(audio_features)
    tokens += [audio_end, tokenizer.eos_token_id]
    return embed(np.array([tokens]), audio_features)


wave, _ = librosa.load("song.flac", sr=16000, mono=True)
track = embed_audio(wave)
query = embed_text("task: search result | query: calm piano")
print(query @ track)

Reproducing the numbers

scripts/ scores the builds on Song Describer, which downloads 3.3 GB on the first run. You only need uv.

cd scripts
uv sync
uv run song_describer.py fp32 fp16 q8 q8:q4 q4

License and credits

Apache-2.0. The graphs come from onnx-community/embeddinggemma-2-ONNX.