sample-identification

Music sample identification: given a query track, retrieve the tracks it samples or is sampled by.

The code and the pretrained checkpoint are released.

Quick start

From a fresh clone to searchable embeddings, using the released checkpoint. No training and no dataset download required.

1. Install

chmod u+x install.sh
./install.sh
conda activate sample-id

This creates a conda environment named sample-id (Python 3.11, ffmpeg<7, and the pip dependencies). See Environment if you would rather install into an environment you already have.

2. Get the checkpoint

The trained weights are on Zenodo: https://zenodo.org/records/22869014 (DOI 10.5281/zenodo.22869014).

mkdir -p checkpoints/sie
BASE=https://zenodo.org/records/22869014/files
curl -L -o checkpoints/sie/best.ckpt   "$BASE/best.ckpt?download=1"    # ~2.4 GB
curl -L -o checkpoints/sie/config.yaml "$BASE/config.yaml?download=1"

best.ckpt and config.yaml are all you need to extract embeddings — the config is required, since the architecture cannot be reconstructed without it. The record also carries hparams.yaml and metrics.csv (see Pretrained checkpoint).

3. Extract embeddings

python main.py predict \
    -c checkpoints/sie/config.yaml \
    --ckpt_path checkpoints/sie/best.ckpt \
    --data.predict_dataset.input_path /path/to/audio \
    --data.predict_dataset.output_dir embeddings/my-tracks

• -c config.yaml is required. The architecture cannot be reconstructed without it, and it has to be passed on the command line — it is a meta-argument that tells Lightning CLI which files to load, so it cannot live inside another config file. The training dataset paths inside it are placeholders; predict discards the training datasets before instantiating anything, so you do not need to fill them in.
• input_path accepts a single .wav file, a directory (scanned recursively for .wav), or a .txt file listing one audio path per line — the list may point at .mp3 and .mp4 as well. Anything else should go through Pre-process audio first.
• Audio is resampled to 16 kHz mono on the fly. Short files are not filtered out unless you set min_duration — see Inference.
• Output: one .npy per input file, shape (n_segments, 1024), FP16, mirroring the input directory tree under output_dir.
• Defaults are 6 s segments with a 3 s hop. To change them, see Inference.

4. Score the embeddings

Retrieval runs on the embeddings written in step 3, not on audio. Each track is used as both a query and a database entry, with its own self-match excluded:

python evaluate.py embeddings/my-tracks /path/to/test.json \
    --output-dir logs/evaluation/quickstart

All evaluation entry points need a ground-truth file — .json for WhoSampled130k, .csv for Sample100 / SamplePairs. To reproduce the published numbers, get the dataset splits and use test.json. Full options and output format are in Evaluation.

Where to go next

You want to	Go to
Know what is in the checkpoint	Pretrained checkpoint
Get the training/evaluation data	Dataset
Convert your own audio to 16 kHz wav	Pre-process audio
Tune segment length, hop, VRAM use	Inference
Change the search, metrics, or output	Evaluation
Train your own model	Train
Run any of this on SLURM	On the cluster

Pretrained checkpoint

SIE — an embedding model that maps 6-second audio segments to 1024-dimensional embeddings in which sampled/sampling track pairs are near neighbours. Trained on WhoSampled130k.

Download: https://zenodo.org/records/22869014 (DOI 10.5281/zenodo.22869014) — see Quick start for the download commands. The weights are released under CC BY-NC-SA 4.0 — not the AGPL that covers this source code.

File	Description
best.ckpt	Lightning checkpoint, best epoch by val-retrieval/comp_metric. Keeps its full training state (optimizer, LR scheduler, loop and callback state), so it can also resume training with --ckpt_path.
config.yaml	Model/data config to pass to main.py. Required — the architecture cannot be reconstructed without it.
hparams.yaml	Hyperparameters as recorded by Lightning.
metrics.csv	Training and validation curves for every epoch.

Training setup:

Backbone	CLEWS (IBN residual blocks, GeM pooling)
Input	VQT, 16 kHz, 258 bins, 36 bins/octave, fmin A0, 25 ms hop
Segment	6 s at inference, 10 s at training (headroom for time-stretch)
Embedding	1024-d, L2-normalised
Loss	Triplet, margin 0.3, hard positive + hard negative mining
Optimiser	AdamW, lr 3e-4, weight decay 0, cosine annealing to 1e-6
Schedule	20 000 steps, batch 394, precision 16-mixed
Augmentation	Additive noise / room IR / microphone IR (p=0.25 each); VQT spectral masking, time-stretch, pitch transpose (p=0.25 each)

Notes:

• The checkpoint stores the network under net._orig_mod. keys, because it was trained with torch.compile. SIELitModule.on_load_checkpoint remaps the keys to whichever layout the live model expects, so it loads correctly whether or not compile is set. setup() only compiles during fit, so predict never compiles at all.
• The dataset paths in config.yaml are placeholders. They are only read by fit/validate; predict discards the training datasets before instantiating anything, so you do not need to fill them in to extract embeddings.

Results table: see the paper.

Environment

chmod u+x install.sh
./install.sh

install.sh creates a conda environment named sample-id with Python 3.11 and ffmpeg<7, then installs this repository into it in editable mode, which pulls in the pip dependencies declared in pyproject.toml. The SLURM scripts under scripts/slurm/ all activate that environment by name.

If you already have Python 3.11+ and ffmpeg<7 available, install the repository directly:

pip install -e ".[wandb]" \
    --extra-index-url https://download.pytorch.org/whl/cu128

The extra index serves the CUDA 12.8 builds of torch/torchaudio; pip cannot read it from pyproject.toml, so it has to be passed on the command line. Drop [wandb] if you will not log to Weights & Biases. requirements.txt lists the same dependencies for tools that do not read pyproject.toml.

Dataset

Training and evaluation use WhoSampled130k — 91,967 annotated sampling relations over 131,539 tracks, partitioned into train, validation and test splits with no track shared between them.

The archive holds annotations and splits only — train.json, val.json, test.json. No audio: tracks are identified by YouTube video id, and the audio has to be retrieved separately. Record format and per-split counts are in the README inside the archive.

Expected on-disk layout:

<root>/splits-<date>/{train,val,test}.json    # annotations and splits
<root>/whosampled-wav-16khz/<id[:2]>/<id>.wav # 16 kHz mono, sharded by id prefix

Audio is sharded by the first two characters of the YouTube id — zhrZ02OcQpw goes to zh/zhrZ02OcQpw.wav, not to the top level. scripts/preprocess_music.py mirrors the input tree, so downloading into <id[:2]>/<id>.<ext> gives the right layout. The dataset paths in configs/default.yaml and scripts/slurm/*.sh are absolute paths on the MTG cluster; point them at your own copy first.

The metadata is for non-commercial scientific research only and may not be redistributed; published results must cite WhoSampled as the source, and the dataset paper (araz2026building). Full terms accompany the dataset release.

Pre-process audio

scripts/preprocess_music.py converts source audio to 16 kHz mono .wav, mirroring the input directory tree.

Example usage: scripts/slurm/preprocess-whosampled.sh

Inference

Inference settings live in configs/inference.yaml. The Quick start passes everything on the command line and takes the defaults for the rest; use the config file when you want to change the extraction settings themselves.

Key	Default	Meaning
model.predict.segment_duration	6.0	Length of each segment, in seconds.
model.predict.hop_duration	3.0	Hop between segment starts, in seconds.
model.predict.layer	proj	Network layer to read embeddings from.
model.predict.chunk_size	1024	Segments per forward pass; raise it to use more VRAM.
model.predict.output_half	true	Store embeddings in FP16.
data.predict_dataset.init_args.min_duration	none (6.0 in configs/inference.yaml)	Skip files shorter than this, in seconds. Unset, nothing is filtered out.

Note that configs/inference.yaml ships with ckpt_path, input_path, output_dir and segment_duration set to ??? — fill all four in before using the file, or override them on the command line, where CLI flags win.

The dataset arguments can be spelled with or without init_args — --data.predict_dataset.init_args.input_path and --data.predict_dataset.input_path both work; the SLURM wrappers use the short form.

The model architecture config (-c checkpoints/sie/config.yaml, or -c logs/.../config.yaml for your own run) must always be passed on the command line, as it is a meta-argument that tells Lightning CLI which files to load and cannot itself live inside a config file.

Option A — edit configs/inference.yaml directly, then run:

python main.py predict -c logs/lightning_logs/version_1/config.yaml -c configs/inference.yaml

Option B — pass values on the command line (overrides the config):

python main.py predict \
    -c logs/lightning_logs/version_1/config.yaml \
    -c configs/inference.yaml \
    --ckpt_path logs/lightning_logs/version_1/checkpoints/last.ckpt \
    --data.predict_dataset.init_args.input_path /path/to/audio_paths.txt \
    --data.predict_dataset.init_args.output_dir artefacts/version_1/emb

Evaluation

Evaluation runs on embeddings already written to disk by predict (see Inference), not on audio. Extract embeddings first, then point one of the evaluation scripts at that directory together with a ground-truth file.

Which script you need depends on the ground-truth format:

Script	Ground truth	Use for
evaluate.py	.json	WhoSampled130k splits (test.json, val.json)
evaluate-old-datasets.py	.csv	Sample100 and SamplePairs
scripts/evaluate-single-query.py	.json	Same as evaluate.py, one query at a time — slower, but avoids holding the full distance matrix in memory

All three take the same two positional arguments — the embedding directory and the ground-truth path — and most of the retrieval options below.

WhoSampled130k

python evaluate.py /path/to/embeddings /path/to/test.json --output-dir logs/evaluation/my-run

Every loaded track is used as both a query and a database entry; each query is scored against all others with its own self-match excluded.

Sample100 / SamplePairs

python evaluate-old-datasets.py /path/to/embeddings /path/to/samples.csv

These datasets separate queries from references, so the database is built one of two ways:

• default — queries and database are disjoint. Queries are the tracks listed as sample_track_id; the database is every other loaded track, which keeps both the reference tracks and the noise/distractor tracks in the search pool.
• --include-queries-in-db — all tracks go into a single pool, and each query is removed from its own result list after searching.

This script also accepts --downsample N, which keeps every N-th segment of each loaded embedding along the time axis (1, the default, keeps all of them).

Options

Common to all three scripts:

Option	Default	Meaning
--similarity-search	NNS	One of NNS, MIPS, MCSS.
--reduction-k	1	Top segment-level distances averaged per track pair. Higher is slower and usually better.
--k	all	Number of top results scored. Omit to score the full database.
--batch-size	32	Query tracks per batch. Not in evaluate-single-query.py.
--chunk-size	none	Database tracks per chunk; set it if the distance matrix does not fit in VRAM. Not in evaluate-single-query.py.
--load-threads	6	Threads used to load .npy embeddings.
--no-gpu	off	Force CPU.
--output-dir	logs/evaluation/<timestamp>	Where results are written.

Output

Each run writes into --output-dir:

• metrics.json — aggregated metrics (M-AP, M-JNAR, M-R1, M-RR, Top1, Top10, Top100, Top1000, and 95% confidence intervals)
• per-query-metrics.json — the same metrics per query track
• args.json — the arguments the run was invoked with
• commit.txt — the git commit the run was produced from

scripts/evaluate-single-query.py writes only metrics.json.

Train

All subcommands (fit, validate, predict) go through main.py.

Logging config is kept separate so it can be composed on top of any run config. configs/logging.yaml sets up CSV logging, checkpointing, and a progress bar. To also enable Weights & Biases, append configs/wandb.yaml after the logging config (order matters — jsonargparse replaces list-valued keys, and wandb.yaml redeclares the full logger list with both CSV and Wandb):

# CSV logging only
python main.py fit -c configs/default.yaml -c configs/logging.yaml

# CSV + Weights & Biases
python main.py fit -c configs/default.yaml -c configs/logging.yaml -c configs/wandb.yaml

Important note: fit must come before -c!

Checkpoints are written to logs/lightning_logs/version_N/checkpoints/:

• checkpoint-best — best checkpoint by val-retrieval/comp_metric
• checkpoint-epoch_N — most recent epoch checkpoint

Resuming Training

Resume model, optimizer, and epoch state with --ckpt_path:

# Without Wandb
python main.py fit -c configs/default.yaml -c configs/logging.yaml \
    --ckpt_path logs/lightning_logs/version_N/checkpoints/last.ckpt

# With Wandb (resumes the same run — id stored in version_N/wandb_run_id.txt)
python main.py fit -c configs/default.yaml -c configs/logging.yaml -c configs/wandb.yaml \
    --ckpt_path logs/lightning_logs/version_N/checkpoints/last.ckpt

Validation

Run a full validation epoch (retrieval + loss) on a pre-trained checkpoint:

python main.py validate \
    -c logs/lightning_logs/version_N/config.yaml \
    --ckpt_path logs/lightning_logs/version_N/checkpoints/last.ckpt

To run only retrieval validation (skip the loss loader):

python main.py validate \
    -c logs/lightning_logs/version_N/config.yaml \
    --ckpt_path logs/lightning_logs/version_N/checkpoints/last.ckpt \
    --data.val_loss_dataset=null

On the cluster

SLURM wrappers live in scripts/slurm/. They activate the sample-id environment and forward any extra arguments to the underlying script:

# WhoSampled130k: evaluate.sh EMB_DIR GT_PATH [extra args...]
sbatch scripts/slurm/evaluate.sh /path/to/embeddings /path/to/test.json --reduction-k 3

# Sample100 / SamplePairs: the ground-truth CSV is inferred from the embedding
# directory name, which must contain "sample100" or "samplepairs".
sbatch scripts/slurm/evaluate-old-datasets.sh /path/to/embeddings/sample100

To go from a checkpoint to metrics on all three datasets in one job — it runs predict for each dataset, then evaluates each:

# CONFIG CKPT OUTPUT_BASE_DIR [dataset_index] [extra args...]
sbatch scripts/slurm/inference-and-evaluate-all-datasets-single-gpu.sh \
    logs/lightning_logs/version_N/config.yaml \
    logs/lightning_logs/version_N/checkpoints/checkpoint-best.ckpt \
    /path/to/embeddings/version_N

Pass a dataset index as the 4th argument (0 sample100, 1 samplepairs, 2 whosampled-test) to process only that one.

Note that the dataset paths inside these wrappers, and the #SBATCH --chdir line at the top of each one, are absolute paths on the MTG cluster; edit them to point at your own copy of the code and data.

Citation

License

Copyright (c) 2026 Music Technology Group, Universitat Pompeu Fabra.

Released under the GNU Affero General Public License v3.0 or later — see LICENSE.

This covers the source code only. The released pretrained checkpoint is licensed separately, under CC BY-NC-SA 4.0: non-commercial use, with attribution, and adapted models shared alike.

The dataset is not covered by either: WhoSampled130k metadata is for non-commercial scientific research only and may not be redistributed — see Dataset for the full terms.