# wav2taste

[![arXiv](https://img.shields.io/badge/arXiv-2607.03296-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2607.03296)
[![HF Dataset](https://img.shields.io/badge/🤗%20Dataset-sonic--seasoning-FFD21E)](https://huggingface.co/datasets/csc-unipd/sonic-seasoning)
[![HF Models](https://img.shields.io/badge/🤗%20Models-wav2taste-FFD21E)](https://huggingface.co/csc-unipd/wav2taste)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE)
[![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue?logo=python&logoColor=white)](pyproject.toml)

Predict the five taste ratings (`sweet`, `bitter`, `salty`, `sour`, `spicy`) of an audio clip from sound alone. This is the code and the trained models behind the paper *Taste-aware music retrieval from audio embeddings* (Spanio & Rodà, CBMI 2026).

![wav2taste architecture: 15 s audio → one or more frozen encoders → gated concat → two-layer MLP head → five taste ratings in [0,1]](docs/architecture.png)

Two ways in:

- **Use a trained model** to score your own audio (no dataset access needed) — see below.
- **Reproduce the paper** by training the swappable-encoder ablation suite on the [`csc-unipd/sonic-seasoning`](https://huggingface.co/datasets/csc-unipd/sonic-seasoning) dataset (needs Hugging Face access) — see [Reproducing the paper](#reproducing-the-paper).

## Use a pretrained taste model

```bash
pip install git+https://github.com/CSCPadova/wav2taste
# or, in this repo:  uv sync
```

```python
from wav2taste import load_taste_model

model = load_taste_model("vggish+mule")     # best fusion; downloads the head + encoders on first use
print(model.predict("song.wav"))
# {'sweet': 0.34, 'bitter': 0.71, 'salty': 0.12, 'sour': 0.58, 'spicy': 0.09}
```

The head is downloaded from [`csc-unipd/wav2taste`](https://huggingface.co/csc-unipd/wav2taste); the frozen audio encoders come from their own upstream repos. Inference needs no access to the private training dataset. Available model names (`wav2taste.PRETRAINED`):

- single encoders: `mfcc`, `clap`, `mert`, `ast`, `encodec`, `vggish`, `hubert`, `panns`, `omar-rq`, `mule`
- gated fusions: `ast+vggish`, `ast+mule`, `vggish+mule` (best), `clap+mule`, `clap+vggish`, `ast+vggish+mule`, `clap+ast+vggish+mule`

> **Weights license:** the trained heads are released under **CC-BY-NC-4.0** (they derive from a non-commercial dataset and from MULE's CC-BY-NC weights). The *code* in this repository is Apache-2.0 (see [License](#license)).

The whole project is built around **swappable audio encoders** so you can run ablations with one CLI flag.

## Architecture


One or more **frozen** audio encoders produce per-encoder embeddings that are concatenated and re-weighted by a learned per-encoder gate; the gated representation feeds a shared two-layer MLP whose sigmoid head outputs a 5-D taste vector in `[0, 1]⁵`.

- **One model, five outputs** (multi-task). Five independent regressors would throw away the strong correlations between tastes (`sweet`/`bitter`, `sweet`/`sour`).
- **Masked MSE** ignores cells where a row didn't get a rating for that taste.
- **Frozen encoders** are the default for ablation runs — we vary the encoder, not its training dynamics. Pass `--no-freeze --no-cache` to fine-tune end-to-end.

## Encoders shipped in the registry

| name        | model                                                 | SR     | embedding |
|-------------|-------------------------------------------------------|-------:|----------:|
| `mfcc`      | hand-crafted MFCC mean+std (no pretraining; baseline) | 22 050 |        80 |
| `clap`      | `laion/clap-htsat-unfused`                            | 48 000 |       768 |
| `mert`      | `m-a-p/MERT-v1-95M`                                   | 24 000 |       768 |
| `ast`       | `MIT/ast-finetuned-audioset-10-10-0.4593`             | 16 000 |       768 |
| `encodec`   | `facebook/encodec_24khz` (pre-quantization latent)    | 24 000 |       128 |
| `vggish`    | Google VGGish (AudioSet, via `torchvggish`)           | 16 000 |       128 |
| `hubert`    | `facebook/hubert-base-ls960` (speech SSL — control)   | 16 000 |       768 |
| `panns`     | PANNs CNN14 (AudioSet, via `panns-inference`)         | 32 000 |     2 048 |
| `omar-rq`   | `mtg-upf/omar-rq-multifeature-25hz-fsq` (layer 6)     | 16 000 |       768 |
| `mule`      | MULE SF-NFNet-F0 (via `mule-torch`, weights [`matteospanio/mule`](https://huggingface.co/matteospanio/mule)) | 16 000 | 1 728 |

(Numbers are nominal — actual `embedding_dim` is read from the model at construction.)

## Reproducing the paper

Everything below trains and evaluates the benchmark from scratch on the private dataset. The exact, cluster-agnostic command sequence behind every paper table and figure is in [`REPRODUCE.md`](REPRODUCE.md); the analysis subcommands live under `wav2taste study`.

### Setup

```bash
# 1. install
uv sync

# 2. authenticate to Hugging Face (the dataset is private)
huggingface-cli login    # or:  export HF_TOKEN=...
```

## Quickstart

```bash
# Precompute frozen-encoder embeddings (one-time per encoder).
uv run wav2taste cache --encoder clap

# Train the head (fast — only the small MLP gets gradients).
uv run wav2taste train --encoder clap --epochs 100

# Evaluate the saved checkpoint on test.
uv run wav2taste eval --checkpoint runs/clap/best.pt --split test
```

## GPU-first inference & ONNX export

For extracting embeddings over a **large** audio corpus, the encoders have a
GPU-native, ONNX-exportable inference path. Feature extraction (mel / fbank /
STFT) is reimplemented in pure torch (`src/wav2taste/frontends/`,
STFT as a conv1d-DFT so it survives ONNX export — `torch.stft` does not), so
`encoder.forward((B, T)) -> (B, D)` is a single batched tensor graph with no
numpy round-trips or per-clip Python loops. This makes the frozen-encoder cache
faster *and* lets the model run under `onnxruntime` with a CUDA execution
provider. Supported encoders: `vggish`, `ast`, `clap`, `mert`.

```bash
uv sync --group onnx                  # adds onnx, onnxruntime, pyarrow
# (on a GPU box install onnxruntime-gpu instead of onnxruntime)

# Export an encoder to ONNX (embeddings only), validating against torch.
uv run wav2taste export --encoder vggish --out vggish.onnx --validate

# Or fold a trained head in to also emit the 5 taste values.
uv run wav2taste export --encoder ast --checkpoint runs/ast/best.pt \
    --emit-taste --out ast_taste.onnx --validate

# Extract embeddings over a directory (or a CSV manifest) of audio, on GPU,
# resumable and streamed to parquet.
uv run wav2taste extract --model vggish.onnx --input /data/songs \
    --glob "**/*.mp3" --out embeddings.parquet --providers cuda,cpu --resume
```

The exported graph takes a **fixed-length** window (`chunk_seconds * sample_rate`
samples; the batch axis is dynamic) and `wav2taste extract` handles variable
length by host-side decode + resample to the encoder's rate, slicing each clip
into windows and mean-pooling the per-window embeddings. Export geometry is
written to a `<model>.onnx.json` sidecar that the extractor reads. The `mule`
encoder is the torch-native port (the `mule-torch` package; see
`encoders/mule_torch.py`), so it joins this path like every other encoder.

## Running an ablation sweep

```bash
for ENC in mfcc clap mert ast encodec vggish hubert panns omar-rq; do
    uv run wav2taste cache --encoder "$ENC"
    uv run wav2taste train --encoder "$ENC" --output-dir "runs/$ENC"
done
```

Each run drops a `runs/<encoder>/summary.json` containing per-target Pearson r / MAE / RMSE on val and test. Compare across encoders by reading those files.

## Paper baseline: per-taste SOTA AST regressors

To anchor the ablation against published reference models, a fifth-of-its-kind baseline runs **five** fine-tuned AST regressors — one per taste — at `csc-unipd/ast-finetuned-{sweet,bitter,salty,sour,spicy}`. Each repo is a single-output `ASTForAudioClassification` (`num_labels=1`) trained as a regressor for one taste dimension. Outputs are concatenated in `TARGETS` order so the result drops into the comparison table next to the multi-task heads in `runs/`.

```bash
uv run wav2taste sota --split test --output-dir runs/sota
```

This shares one feature-extraction pass across the five heads and writes `runs/sota/summary.json` with the same `test_metrics` shape as a training run.

## Paper-ready comparison table

After training the encoders and running the SOTA baseline, aggregate every `runs/*/summary.json` into a single side-by-side table:

```bash
uv run wav2taste compare \
    --order sota mfcc clap mert ast encodec vggish hubert panns omar-rq \
    --output runs/comparison.md \
    --latex  runs/comparison.tex
```

This emits one Markdown and one LaTeX block per metric (Pearson r, MAE, RMSE) with the best entry per column bolded — drop-in for the experiments section of a paper.

End-to-end:

```bash
for ENC in mfcc clap mert ast encodec vggish hubert panns omar-rq; do
    uv run wav2taste cache --encoder "$ENC"
    uv run wav2taste train --encoder "$ENC" --output-dir "runs/$ENC"
done
uv run wav2taste sota --output-dir runs/sota
uv run wav2taste compare --order sota mfcc clap mert ast encodec vggish hubert panns omar-rq
```

## Paper-strengthening studies

The base encoder table is only the first step. The repository also ships a
study CLI for the follow-up analyses needed for a stronger paper narrative.

```bash
# Stability: repeated seeds for one encoder / training regime.
uv run wav2taste study seed-sweep --encoder ast --output-dir runs/ast_seed_sweep

# Isolate the benefit of multi-task learning.
uv run wav2taste study single-task --encoder ast --output-dir runs/ast_single_task

# Compare frozen, LoRA, and full tuning.
uv run wav2taste study adaptation --encoders ast --output-dir runs/adaptation_study

# Fuse complementary encoders on cached embeddings.
uv run wav2taste study fusion --encoders ast vggish --output-dir runs/fusion_ast_vggish

# Probe sparse-target failure modes.
uv run wav2taste study imbalance --encoders clap omar-rq --output-dir runs/imbalance_study
```

To analyze source shift or run paired significance tests, first export raw
predictions from any checkpoint or from the SOTA baseline:

```bash
uv run wav2taste eval --checkpoint runs/ast/best.pt --save-predictions runs/ast/test_predictions.npz
uv run wav2taste sota --save-predictions runs/sota/test_predictions.npz

uv run wav2taste study source-eval \
    --predictions runs/ast/test_predictions.npz runs/sota/test_predictions.npz \
    --output-dir runs/source_eval

uv run wav2taste study significance \
    --first runs/ast/test_predictions.npz \
    --second runs/sota/test_predictions.npz \
    --output-dir runs/significance_ast_vs_sota
```

## Commands reference

One command, `wav2taste <subcommand>` (drop the `uv run` prefix once installed as a package). The full study suite is reproducible end-to-end from [`REPRODUCE.md`](REPRODUCE.md).

### Core entry points

| command | what it does |
|---|---|
| `wav2taste cache --encoder <name>` | Precompute and cache frozen-encoder embeddings for one encoder. |
| `wav2taste train --encoder <name>` | Train the MLP head against the cached embeddings. |
| `wav2taste eval --checkpoint <best.pt>` | Evaluate a saved checkpoint and dump test predictions. |
| `wav2taste sota` | Reference baseline: load + evaluate the five external `csc-unipd/ast-finetuned-{taste}` regressors. |
| `wav2taste compare` | Aggregate `runs/*/summary.json` into a paper-ready Markdown + LaTeX comparison table. |

### Study subcommands (`wav2taste study <subcommand>`)

| subcommand | what it does |
|---|---|
| `seed-sweep` | Repeat one configuration over multiple seeds and report mean ± std. |
| `single-task` | Train one independent head per taste and merge predictions. |
| `source-eval` | Compare saved prediction files broken down by source (perceptual-validation / tasty-musicgen / annotated-corpus). |
| `adaptation` | Frozen vs. LoRA vs. full fine-tuning sweep. |
| `fusion` | Train a gated late-fusion head on multiple cached encoders. |
| `imbalance` | Probe sparse-target failure (multi-task / weighted / single-task variants). |
| `significance` | Paired bootstrap difference between two prediction files. |
| `human-ceiling` | Compare model Pearson r to per-source inter-rater r ceiling. |
| `permutation-test` | Per-target label-shuffle permutation p-values. |
| `retrieval` | Retrieve test items by predicted taste profile (Precision@k, NDCG@k). |
| `cka` | Pairwise linear CKA between cached encoder embeddings. |
| `psychoacoustic-probe` | Ridge probe per encoder onto librosa psychoacoustic descriptors. |
| `kfold` | k-fold CV on train+val with the test split held out. |
| `saliency` | Hand-rolled integrated gradients on AST mel-spec inputs. |
| `crossmodal-verify` | Evaluate a checkpoint on synthetic literature-grounded stimuli. |

Run `wav2taste study <subcommand> --help` for the full argument list (the table above is a subset — `wav2taste study --help` lists them all). Each writes `runs/<name>/summary.json` and `runs/<name>/report.md`; many also save `test_predictions.npz` for downstream comparison.

## Adding a new encoder

1. Create `src/wav2taste/encoders/myenc.py`:

   ```python
   from torch import Tensor
   from wav2taste.encoders.base import AudioEncoder
   from wav2taste.encoders.registry import register

   @register("myenc")
   class MyEncoder(AudioEncoder):
       def __init__(self, ...):
           super().__init__()
           ...
       @property
       def embedding_dim(self) -> int: return ...
       @property
       def sample_rate(self) -> int: return ...
       def forward(self, waveform: Tensor) -> Tensor: ...
   ```

2. Add an import in `encoders/__init__.py` (inside the `_try_import` block).
3. Run `uv run wav2taste cache --encoder myenc` and you're set.

## License

- **Code:** Apache-2.0 (see [`LICENSE`](LICENSE)).
- **Trained heads** ([`csc-unipd/wav2taste`](https://huggingface.co/csc-unipd/wav2taste)): CC-BY-NC-4.0 — they derive from the non-commercial training corpus and from MULE's CC-BY-NC weights, so they are for non-commercial research use.
- The frozen upstream encoders and the training dataset keep their own licenses.

If you use this work, please cite the CBMI 2026 paper:

```bibtex
@inproceedings{spanio2026taste,
  title     = {Taste-aware music retrieval from audio embeddings},
  author    = {Spanio, Matteo and Rod{\`a}, Antonio},
  booktitle = {Proceedings of the International Conference on Content-Based Multimedia Indexing (CBMI)},
  year      = {2026},
}
```

## Handling missing ratings

The dataset is partially-rated: ~257 annotated-corpus rows lack `hot/cold/emotions`, only 78 rows have `spicy`, etc. The training loop uses **masked MSE**, so each row contributes loss only on the columns it was actually rated on. No imputation, no row dropping.

If a target is much sparser than the others (`spicy` has ~4× fewer training rows than `sweet`), pass `--balance-targets` to upweight it inversely to its coverage. By default this is off — start without it and turn it on if `spicy` Pearson r lags far behind the other tastes.