Views
No views yet
pip install .pip install -e .[train]1import torch
2from omar_rq import get_model
3
4# Embedding extraction example
5x = torch.randn(1, 16000 * 4).cpu() # Example: 4 seconds of mono audio at 16kHz
6
7# Load a specific model, e.g., "mtg-upf/omar-rq-multifeature-25hz-fsq"
8model_id = "mtg-upf/omar-rq-multifeature-25hz-fsq"
9model = get_model(model_id=model_id, device="cpu") # Use "cuda" if a GPU is available
10
11# Extract embeddings from layer 6
12embeddings = model.extract_embeddings(x, layers=[6])
13
14# Use the `model.eps` field to compute timestamps for the extracted embeddings
15timestamps = torch.arange(embeddings.shape[2]) / model.eps
16
17print(f"Extracted embeddings shape: {embeddings.shape}")
18print(f"First 5 timestamps: {timestamps[:5]}")get_model reference:Returns an OMAR-RQ Module from the provided model_id or config_file.
Args:
model_id (str): Hugging Face's Model ID or local path to the model
config_file (Path): Path to the model config of a trained model.
device (str): Device to use for the model. Defaults to "cpu".
quantization_targets (bool): If True, it will create the quantization
targets for SSL pre-training of the model. Defaults to False.
Output:
module: The model from the provided config file.
Module usage:
Args:
audio (torch.Tensor): 2D mono audio tensor (B, T'). Where B is
the batch size and T' is the number of samples.
layers (set): Set of layer indices to extract embeddings from.
By default, it extracts embeddings from the last layer (logits).
Output:
torch.Tensor: Extracted embeddings. The output tensor has shape
(L, B, T, C,) where L = len(layers), B is the batch size, T is
the number of output timestamps, and C = embedding dimension.extract_embeddings reference:Extract embeddings from an input audio batch.
Args:
audio (torch.Tensor): 2D mono audio tensor (B, T'). Where B is
the batch size and T' is the number of samples.
layers (set): Set of layer indices to extract embeddings from.
By default, it extracts embeddings from the last layer (logits).
Output:
torch.Tensor: Extracted embeddings. The output tensor has shape
(L, B, T, C,) where L = len(layers), B is the batch size, T is
the number of output timestamps, and C = embedding dimension.| Model | Input | Rate | Tagging | Difficulty | Pitch | Chord | Beat | Structure | Hugging Face Model ID |
|---|---|---|---|---|---|---|---|---|---|
| Hz | mAP | MSE | acc. | acc. | F1 | acc. | |||
| base | mel | 15.63 | .482 | 1.65 | .892 | .657 | .783 | .647 | mtg-upf/omar-rq-base |
| multicodebook | mel | 15.63 | .488 | 1.66 | .897 | .675 | .775 | .639 | mtg-upf/omar-rq-multicodebook |
| multifeature | audio | 18.75 | .467 | 1.76 | .938 | .734 | .833 | .623 | mtg-upf/omar-rq-multifeature |
| multifeature-25hz | audio | 25 | .463 | 1.79 | .932 | .728 | .848 | .628 | mtg-upf/omar-rq-multifeature-25hz |
| multifeature-25hz-fsq | audio | 25 | .463 | 1.71 | .940 | .749 | .855 | .628 | mtg-upf/omar-rq-multifeature-25hz-fsq |
1@article {alonso2025omarrq,
2 title={OMAR-RQ: Open Music Audio Representation Model Trained with Multi-Feature Masked Token Prediction},
3 author={Alonso-Jim\'enez, Pablo and Ramoneda, Pedro and Araz, R. Oguz and Poltronieri, Andrea and Bogdanov, Dmitry},
4 journal={arXiv preprint arXiv:2507.03482},
5 year={2025}
6}