Initial release

.gitattributes

@@ -1,35 +1,37 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+*.onnx_data filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,141 @@
+---
+license: apache-2.0
+language:
+  - en
+library_name: transformers.js
+pipeline_tag: feature-extraction
+base_model: jaisawhney/SonicEmbed
+tags:
+  - onnx
+  - audio
+  - music
+  - music-information-retrieval
+  - text-to-audio-retrieval
+  - embeddings
+---
+
+# SonicEmbed ONNX
+
+ONNX builds of [SonicEmbed](https://huggingface.co/jaisawhney/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](https://huggingface.co/jaisawhney/SonicEmbed) | 2,307 MB | 40.4              | 36.7           | 9.5              | 16                       |
+
+Scores are text-to-music search on Song Describer.
+
+- **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.
+
+## Usage
+
+### Transformers.js
+
+```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
+
+```bash
+pip install onnxruntime transformers huggingface_hub librosa
+```
+
+```python
+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/*{BUILD}.onnx*", "*.json"],
+)
+
+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](https://docs.astral.sh/uv/getting-started/installation/).
+
+```bash
+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](https://huggingface.co/onnx-community/embeddinggemma-2-ONNX).

chat_template.jinja

@@ -0,0 +1 @@
+{%- for msg in messages if msg.get('role') == 'system' -%}{%- if msg.get('content') is string -%}{{ msg['content'] }}{%- else -%}{%- for item in msg['content'] if item.get('type') == 'text' -%}{{ item['text'] }}{%- endfor -%}{%- endif -%}{%- endfor -%}{%- for msg in messages if msg.get('role') != 'system' -%}{%- if msg.get('content') is string -%}{{ msg['content'] }}{%- else -%}{%- set existing_text = msg['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | join -%}{%- set has_manual_placeholders = ('<|image|>' in existing_text) or ('<|video|>' in existing_text) or ('<|audio|>' in existing_text) -%}{%- for item in msg['content'] -%}{%- if item.get('type') == 'text' -%}{{ item['text'] }}{%- elif not has_manual_placeholders and item.get('type') == 'image' -%}<|image|>{%- elif not has_manual_placeholders and item.get('type') == 'video' -%}<|video|>{%- elif not has_manual_placeholders and item.get('type') == 'audio' -%}<|audio|>{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}

config.json

@@ -0,0 +1,130 @@
+{
+  "architectures": [
+    "EmbeddingGemma2Model"
+  ],
+  "audio_config": {
+    "attention_chunk_size": 12,
+    "attention_context_left": 13,
+    "attention_context_right": 0,
+    "attention_invalid_logits_value": -1000000000.0,
+    "attention_logit_cap": 50.0,
+    "conv_kernel_size": 5,
+    "dtype": "float32",
+    "gradient_clipping": 10000000000.0,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "model_type": "gemma4_audio",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 12,
+    "output_proj_dims": 1536,
+    "residual_weight": 0.5,
+    "rms_norm_eps": 1e-06,
+    "subsampling_conv_channels": [
+      128,
+      32
+    ],
+    "use_clipped_linears": true
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "dtype": "float32",
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "model_type": "embedding_gemma2",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 2,
+    "dtype": "float32",
+    "embedding_dim": 768,
+    "eos_token_id": 1,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 512,
+    "hidden_size_per_layer_input": 512,
+    "initializer_range": 0.02,
+    "intermediate_size": 2048,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "embedding_gemma2_text",
+    "num_attention_heads": 4,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "pad_token_id": 0,
+    "per_layer_config": {
+      "05": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "11": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "17": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      },
+      "23": {
+        "head_dim": 512,
+        "num_key_value_heads": 1
+      }
+    },
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "rope_theta": 1000000.0,
+        "rope_type": "default"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 512,
+    "vocab_size": 262144
+  },
+  "transformers_version": "5.20.2",
+  "video_token_id": 258884,
+  "vision_soft_tokens_per_image": 280,
+  "transformers.js_config": {
+    "use_external_data_format": {
+      "audio_encoder.onnx": 1,
+      "model.onnx": 1,
+      "audio_encoder_fp16.onnx": 1,
+      "model_fp16.onnx": 1,
+      "audio_encoder_quantized.onnx": 1,
+      "model_quantized.onnx": 1,
+      "audio_encoder_q4.onnx": 1,
+      "model_q4.onnx": 1
+    }
+  }
+}

config_sentence_transformers.json

@@ -0,0 +1,32 @@
+{
+  "__version__": {
+    "pytorch": "2.11.0+cu128",
+    "sentence_transformers": "6.1.0",
+    "transformers": "5.20.2"
+  },
+  "default_prompt_name": null,
+  "model_type": "SentenceTransformer",
+  "prompts": {
+    "BitextMining": "task: search result | query: ",
+    "Classification": "task: classification | query: ",
+    "Clustering": "task: clustering | query: ",
+    "CodeRetrieval": "task: code retrieval | query: ",
+    "Document": "title: none | text: ",
+    "FactChecking": "task: fact checking | query: ",
+    "InstructionRetrieval": "task: code retrieval | query: ",
+    "MultilabelClassification": "task: classification | query: ",
+    "PairClassification": "task: sentence similarity | query: ",
+    "QuestionAnswering": "task: question answering | query: ",
+    "Reranking": "task: search result | query: ",
+    "Retrieval": "task: search result | query: ",
+    "Retrieval-document": "title: none | text: ",
+    "Retrieval-query": "task: search result | query: ",
+    "STS": "task: sentence similarity | query: ",
+    "SearchQuery": "task: search result | query: ",
+    "SentenceSimilarity": "task: sentence similarity | query: ",
+    "Summarization": "task: sentence similarity | query: ",
+    "document": "title: none | text: ",
+    "query": "task: search result | query: "
+  },
+  "similarity_fn_name": "cosine"
+}

onnx/audio_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c2cd75556bd606e591c370c8b5edd3c708b0f6d6fc3a11c8a7c415df0181989
+size 204551

onnx/audio_encoder.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7871094ac28b10567a2e6f8cf5e826a0196d46ad9ef034b395c4f6ed87699c6
+size 1172750848

onnx/audio_encoder_fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:644e2709393d15620eb6f62c032632479d48ab7431163d0335fba10b172178c6
+size 201983

onnx/audio_encoder_fp16.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abcf6d2943303c29bdb2ecce22e731246100c5dabb5035ed7389f51c84785ea2
+size 586375424

onnx/audio_encoder_q4.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abf0d00ab23205026fdd4010e36b39b7c9ee17a480d601cd95a81379d35fd047
+size 246624

onnx/audio_encoder_q4.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c240c8f930c12b07a514dc18709692969216e5a6805040669ba35bf6a165c41
+size 189075968

onnx/audio_encoder_quantized.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3d4b59dc4ca20fdaaed962f572bf22b00e2a23f9376a661553abde9b7066b44
+size 250411

onnx/audio_encoder_quantized.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1066a923825875ce7457bd0a7253cac86ccd01d2419070e4d1729ac64e87263f
+size 340058624

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc47de15f81208a5c99e5ab10f746d5e33b51ea228b7dc0bef9c133a94f1c1c3
+size 421421

onnx/model.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a5c3dd3fc284af962104f19dff276fce49ace5b5a052c791c69b09068fda419
+size 1084170240

onnx/model_fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:578e3e3b68dd394cfab8268b914b6c3ccd2b9748e214c1ebe048e73716a51c7f
+size 425780

onnx/model_fp16.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6db39f86c7a0d811a524222cbfd7f5144f0d26588269927528f6be845d7438a
+size 542085120

onnx/model_q4.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc27af510b44b555f1909bd108ddfabac9a5877ca679b83bead8c9eb46787733
+size 492114

onnx/model_q4.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b63420bfe7c8f82ba412572b5ee4a50b5060c7033df5fe08987320d374d7d7aa
+size 174028800

onnx/model_quantized.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61cc82205da352f94ae362e092254b1d5c00e44f422e2a94ece5121d06a5f211
+size 496200

onnx/model_quantized.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:139048d4e312feeda8824cde9b9ba2469d25f633b659adbd2c8f1587a5cc75b4
+size 313724928

preprocessor_config.json

@@ -0,0 +1,22 @@
+{
+  "dither": 0.0,
+  "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+  "feature_size": 128,
+  "fft_length": 512,
+  "fft_overdrive": false,
+  "frame_length": 320,
+  "hop_length": 160,
+  "input_scale_factor": 1.0,
+  "max_frequency": 8000.0,
+  "mel_floor": 0.001,
+  "min_frequency": 0.0,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "per_bin_mean": null,
+  "per_bin_stddev": null,
+  "preemphasis": 0.0,
+  "preemphasis_htk_flavor": true,
+  "return_attention_mask": true,
+  "sampling_rate": 16000,
+  "processor_class": "EmbeddingGemma2Processor"
+}

processor_config.json

@@ -0,0 +1,78 @@
+{
+  "audio_ms_per_token": 40,
+  "audio_seq_length": 280,
+  "feature_extractor": {
+    "dither": 0.0,
+    "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+    "feature_size": 128,
+    "fft_length": 512,
+    "fft_overdrive": false,
+    "frame_length": 320,
+    "hop_length": 160,
+    "input_scale_factor": 1.0,
+    "max_frequency": 8000.0,
+    "mel_floor": 0.001,
+    "min_frequency": 0.0,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "per_bin_mean": null,
+    "per_bin_stddev": null,
+    "preemphasis": 0.0,
+    "preemphasis_htk_flavor": true,
+    "return_attention_mask": true,
+    "sampling_rate": 16000
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_processor_type": "Gemma4ImageProcessor",
+    "image_seq_length": 280,
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 280,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098
+  },
+  "image_seq_length": 280,
+  "processor_class": "EmbeddingGemma2Processor",
+  "video_processor": {
+    "add_timestamps": false,
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "fps": 1,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_frames": 32,
+    "max_soft_tokens": 140,
+    "overflow_strategy": "uniform",
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "video_processor_type": "EmbeddingGemma2VideoProcessor"
+  }
+}

scripts/pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "sonicembed-onnx"
+version = "0.1.0"
+description = "Scores the SonicEmbed ONNX files"
+requires-python = ">=3.11,<3.13"
+dependencies = [
+    "onnxruntime==1.31.1",
+    "huggingface-hub>=1.0",
+    "numpy>=2",
+    "sentence-transformers>=6.1",
+    "librosa>=0.11",
+    "pyarrow>=20",
+    "pandas>=3.1.1",
+    "tqdm",
+    "pillow>=11",
+    "torchvision>=0.23",
+]

scripts/song_describer.py

@@ -0,0 +1,127 @@
+"""Text-to-music search on Song Describer with the ONNX builds: each of the 1,106 captions looks for its track among all 706.
+
+Tracks are embedded two ways: one 10 s clip from the middle, and the average of three 10 s clips.
+AUDIO:TEXT mixes two builds: q8:q4 uses q8 for the audio encoder and q4 for the text model.
+
+    uv run song_describer.py fp32 fp16 q8 q8:q4 q4
+"""
+import argparse
+import io
+from pathlib import Path
+
+import librosa
+import numpy as np
+import onnxruntime as ort
+import pyarrow.parquet as pq
+from huggingface_hub import snapshot_download
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+
+MODEL = "jaisawhney/SonicEmbed"
+MODEL_REVISION = "16c73997cfc80d00a71f06b1fa821465aa2c7e2d"
+DATASET = "dukesun99/SongDescriber-T2A"  # the official 2-minute excerpts, byte for byte
+REVISION = "c58ad9b08343e56ce412d4da41accd15bffdbd6d"
+SR = 16000
+FILE_SUFFIXES = {"fp32": "", "fp16": "_fp16", "q8": "_quantized", "q4": "_q4"}
+NO_FEATURES = np.zeros((0, 512), "float32")
+
+
+def clips(wave, sr=SR):
+    size = 10 * sr
+    starts = [min(max(int(position * len(wave)) - size // 2, 0), len(wave) - size) for position in (0.25, 0.5, 0.75)]
+    return [wave[start:start + size] for start in starts]
+
+
+def load_captions(data):
+    print("Downloading Song Describer (3.3 GB, first run only)")
+    snapshot_download(DATASET, repo_type="dataset", revision=REVISION, local_dir=data)
+    queries = pq.read_table(data / "queries").to_pandas()
+    qrels = pq.read_table(data / "qrels").to_pandas()
+    return queries.merge(qrels, left_on="id", right_on="query-id")
+
+
+def load_tracks(data, sr=SR):
+    shards = sorted((data / "corpus").glob("*.parquet"))
+    with tqdm(total=sum(pq.ParquetFile(shard).metadata.num_rows for shard in shards), desc="tracks") as bar:
+        for shard in shards:
+            for row in pq.read_table(shard).to_pylist():
+                wave, _ = librosa.load(io.BytesIO(row["audio"]["bytes"]), sr=sr, mono=True)
+                yield row["id"], wave
+                bar.update()
+
+
+def recall_table(text, audio, caption_tracks, tracks):
+    sims = text @ audio.T
+    correct = sims[np.arange(len(text)), np.searchsorted(tracks, caption_tracks)]
+    ranks = (sims > correct[:, None]).sum(1) + 1
+    return {"R@1": 100 * (ranks <= 1).mean(), "R@5": 100 * (ranks <= 5).mean(),
+            "R@10": 100 * (ranks <= 10).mean(), "MedR": np.median(ranks), "MRR": (1 / ranks).mean()}
+
+
+def print_results(name, text, views, caption_tracks, tracks):
+    print(f"\n{name}: {len(text)} captions, {len(tracks)} tracks\n")
+    print("| Audio | R@1 | R@5 | R@10 | MedR | MRR |\n|---|---|---|---|---|---|")
+    for view, audio in views.items():
+        scores = recall_table(text, audio, caption_tracks, tracks)
+        print(f"| {view} | {scores['R@1']:.1f} | {scores['R@5']:.1f} | {scores['R@10']:.1f} | {scores['MedR']:g} | {scores['MRR']:.4f} |")
+
+
+def open_session(path):
+    options = ort.SessionOptions()
+    options.log_severity_level = 3
+    return ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"])
+
+
+def embed_text(text_model, inputs, audio_features=NO_FEATURES):
+    output = text_model.run(["sentence_embedding"], {
+        "input_ids": inputs["input_ids"].numpy(), "attention_mask": inputs["attention_mask"].numpy(),
+        "image_features": NO_FEATURES, "video_features": NO_FEATURES, "audio_features": audio_features})[0][0]
+    return output / np.linalg.norm(output)
+
+
+def embed_clip(audio_encoder, text_model, inputs):
+    features = audio_encoder.run(None, {
+        "input_features": inputs["input_features"].numpy().astype("float32"),
+        "input_features_mask": inputs["input_features_mask"].numpy().astype(bool)})[0]
+    return embed_text(text_model, inputs, features)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("builds", nargs="+", help="fp32, fp16, q8, q4 or AUDIO:TEXT")
+    parser.add_argument("--onnx", type=Path, default=Path("../onnx"))
+    parser.add_argument("--data", type=Path, default=Path("data/song_describer"))
+    args = parser.parse_args()
+
+    captions = load_captions(args.data)
+    model = SentenceTransformer(MODEL, revision=MODEL_REVISION, device="cpu", config_kwargs={"vision_config": None})
+
+    tracks, clip_inputs = [], []
+    for track, wave in load_tracks(args.data):
+        tracks.append(track)
+        clip_inputs += [model.preprocess([{"array": clip, "sampling_rate": SR}]) for clip in clips(wave)]
+    caption_inputs = [model.preprocess([caption], prompt=model.prompts["query"]) for caption in captions.text]
+
+    order = np.argsort(tracks)
+    tracks = np.array(tracks)[order]
+    caption_tracks = captions["corpus-id"].to_numpy()
+
+    for build in args.builds:
+        audio_build, text_build = build, build
+        if ":" in build:
+            audio_build, text_build = build.split(":")
+        audio_encoder = open_session(args.onnx / f"audio_encoder{FILE_SUFFIXES[audio_build]}.onnx")
+        text_model = open_session(args.onnx / f"model{FILE_SUFFIXES[text_build]}.onnx")
+
+        emb = np.stack([embed_clip(audio_encoder, text_model, inputs) for inputs in tqdm(clip_inputs, desc=f"{build} audio")])
+        text = np.stack([embed_text(text_model, inputs) for inputs in tqdm(caption_inputs, desc=f"{build} text")])
+
+        emb = emb.reshape(len(tracks), 3, -1)
+        middle, average = emb[:, 1][order], emb.mean(1)[order]
+        average = average / np.linalg.norm(average, axis=-1, keepdims=True)
+
+        print_results(build, text, {"one 10 s clip": middle, "three 10 s clips": average}, caption_tracks, tracks)
+
+
+if __name__ == "__main__":
+    main()

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d777ef5bdc1aa36227abdfb77c3e49e7b9c892d16e1b6bda41c393504828be4
+size 32170510

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>",
+    "video_token": "<|video|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "right",
+  "processor_class": "EmbeddingGemma2Processor",
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>",
+  "video_token": "<|video|>"
+}