init

MERT-v1-95M/README.md

@@ -0,0 +1,121 @@
+---
+license: mit
+inference: false
+tags:
+- music
+---
+
+# Introduction to our series work
+
+The development log of our Music Audio Pre-training (m-a-p) model family:
+- 17/03/2023: we release two advanced music understanding models, [MERT-v1-95M](https://huggingface.co/m-a-p/MERT-v1-95M) and [MERT-v1-330M](https://huggingface.co/m-a-p/MERT-v1-330M) , trained with new paradigm and dataset. They outperform the previous models and can better generalize to more tasks.
+- 14/03/2023: we retrained the MERT-v0 model with open-source-only music dataset [MERT-v0-public](https://huggingface.co/m-a-p/MERT-v0-public)
+- 29/12/2022: a music understanding model [MERT-v0](https://huggingface.co/m-a-p/MERT-v0) trained with **MLM** paradigm, which performs better at downstream tasks.
+- 29/10/2022: a pre-trained MIR model [music2vec](https://huggingface.co/m-a-p/music2vec-v1) trained with **BYOL** paradigm.
+
+
+
+Here is a table for quick model pick-up:
+
+| Name                                                         | Pre-train Paradigm | Training Data (hour) | Pre-train Context   (second) | Model Size | Transformer Layer-Dimension | Feature Rate | Sample Rate | Release Date |
+| ------------------------------------------------------------ | ------------------ | -------------------- | ---------------------------- | ---------- | --------------------------- | ------------ | ----------- | ------------ |
+| [MERT-v1-330M](https://huggingface.co/m-a-p/MERT-v1-330M)    | MLM                | 160K                 | 5                            | 330M       | 24-1024                     | 75 Hz        | 24K Hz      | 17/03/2023   |
+| [MERT-v1-95M](https://huggingface.co/m-a-p/MERT-v1-95M)      | MLM                | 20K                  | 5                            | 95M        | 12-768                      | 75 Hz        | 24K Hz      | 17/03/2023   |
+| [MERT-v0-public](https://huggingface.co/m-a-p/MERT-v0-public) | MLM                | 900                  | 5                            | 95M        | 12-768                      | 50 Hz        | 16K Hz      | 14/03/2023   |
+| [MERT-v0](https://huggingface.co/m-a-p/MERT-v0)              | MLM                | 1000                 | 5                            | 95 M       | 12-768                      | 50 Hz        | 16K Hz      | 29/12/2022   |
+| [music2vec-v1](https://huggingface.co/m-a-p/music2vec-v1)    | BYOL               | 1000                 | 30                           | 95 M       | 12-768                      | 50 Hz        | 16K Hz      | 30/10/2022   |
+
+## Explanation
+
+The m-a-p models share the similar model architecture and the most distinguished difference is the paradigm in used pre-training. Other than that, there are several nuance technical configuration needs to know before using:
+
+- **Model Size**: the number of parameters that would be loaded to memory. Please select the appropriate size fitting your hardware.
+- **Transformer Layer-Dimension**: The number of transformer layers and the corresponding feature dimensions can be outputted from our model. This is marked out because features extracted by **different layers could have various performance depending on tasks**.
+- **Feature Rate**: Given a 1-second audio input, the number of features output by the model.
+- **Sample Rate**: The frequency of audio that the model is trained with.
+
+
+
+# Introduction to MERT-v1
+
+Compared to MERT-v0, we introduce multiple new things in the MERT-v1 pre-training:
+
+- Change the pseudo labels to 8 codebooks from [encodec](https://github.com/facebookresearch/encodec), which potentially has higher quality and empower our model to support music generation.
+- MLM prediction with in-batch noise mixture.
+- Train with higher audio frequency (24K Hz).
+- Train with more audio data (up to 160 thousands of hours).
+- More available model sizes 95M and 330M.
+
+
+
+More details will be written in our coming-soon paper.
+
+
+
+# Model Usage
+
+```python
+# from transformers import Wav2Vec2Processor
+from transformers import Wav2Vec2FeatureExtractor
+from transformers import AutoModel
+import torch
+from torch import nn
+import torchaudio.transforms as T
+from datasets import load_dataset
+
+
+# loading our model weights
+model = AutoModel.from_pretrained("m-a-p/MERT-v1-95M", trust_remote_code=True)
+# loading the corresponding preprocessor config
+processor = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v1-95M",trust_remote_code=True)
+
+# load demo audio and set processor
+dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+dataset = dataset.sort("id")
+sampling_rate = dataset.features["audio"].sampling_rate
+
+resample_rate = processor.sampling_rate
+# make sure the sample_rate aligned
+if resample_rate != sampling_rate:
+    print(f'setting rate from {sampling_rate} to {resample_rate}')
+    resampler = T.Resample(sampling_rate, resample_rate)
+else:
+    resampler = None
+
+# audio file is decoded on the fly
+if resampler is None:
+	input_audio = dataset[0]["audio"]["array"]
+else:
+  input_audio = resampler(torch.from_numpy(dataset[0]["audio"]["array"]))
+  
+inputs = processor(input_audio, sampling_rate=resample_rate, return_tensors="pt")
+with torch.no_grad():
+    outputs = model(**inputs, output_hidden_states=True)
+
+# take a look at the output shape, there are 13 layers of representation
+# each layer performs differently in different downstream tasks, you should choose empirically
+all_layer_hidden_states = torch.stack(outputs.hidden_states).squeeze()
+print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim]
+
+# for utterance level classification tasks, you can simply reduce the representation in time
+time_reduced_hidden_states = all_layer_hidden_states.mean(-2)
+print(time_reduced_hidden_states.shape) # [13, 768]
+
+# you can even use a learnable weighted average representation
+aggregator = nn.Conv1d(in_channels=13, out_channels=1, kernel_size=1)
+weighted_avg_hidden_states = aggregator(time_reduced_hidden_states.unsqueeze(0)).squeeze()
+print(weighted_avg_hidden_states.shape) # [768]
+```
+
+
+
+# Citation
+
+```shell
+@article{li2022large,
+  title={Large-Scale Pretrained Model for Self-Supervised Music Audio Representation Learning},
+  author={Li, Yizhi and Yuan, Ruibin and Zhang, Ge and Ma, Yinghao and Lin, Chenghua and Chen, Xingran and Ragni, Anton and Yin, Hanzhi and Hu, Zhijie and He, Haoyu and others},
+  year={2022}
+}
+
+```

MERT-v1-95M/config.json

@@ -0,0 +1,85 @@
+{
+  "_name_or_path": "m-a-p/MERT-v1-95M",
+  "activation_dropout": 0.0,
+  "apply_spec_augment": true,
+  "architectures": [
+    "MERTModel"
+  ],
+  "attention_dropout": 0.1,
+  "attention_relax": -1.0,
+  "auto_map": {
+    "AutoConfig": "configuration_MERT.MERTConfig",
+    "AutoModel": "modeling_MERT.MERTModel"
+  },
+  "bos_token_id": 1,
+  "classifier_proj_size": 256,
+  "conv_bias": false,
+  "conv_dim": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "conv_kernel": [
+    10,
+    3,
+    3,
+    3,
+    3,
+    2,
+    2
+  ],
+  "conv_stride": [
+    5,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2
+  ],
+  "ctc_loss_reduction": "sum",
+  "ctc_zero_infinity": false,
+  "deepnorm": false,
+  "do_stable_layer_norm": false,
+  "eos_token_id": 2,
+  "feat_extract_activation": "gelu",
+  "feat_extract_dropout": 0.0,
+  "feat_extract_norm": "group",
+  "feat_proj_dropout": 0.1,
+  "feat_proj_layer_norm": true,
+  "feature_extractor_cqt": false,
+  "feature_extractor_cqt_bins": 336,
+  "final_dropout": 0.1,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout": 0.1,
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "layerdrop": 0.05,
+  "mask_feature_length": 10,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_prob": 0.05,
+  "model_type": "mert_model",
+  "num_attention_heads": 12,
+  "num_conv_pos_embedding_groups": 16,
+  "num_conv_pos_embeddings": 128,
+  "num_feat_extract_layers": 7,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "sample_rate": 24000,
+  "tokenizer_class": "Wav2Vec2CTCTokenizer",
+  "torch_dtype": "float32",
+  "transformers_version": "4.24.0",
+  "use_weighted_layer_sum": false,
+  "vocab_size": 32
+}

MERT-v1-95M/configuration_MERT.py

@@ -0,0 +1,138 @@
+"""
+MERT model configuration
+"""
+
+import functools
+import operator
+
+# from ...configuration_utils import PretrainedConfig
+# from ...utils import logging
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+# TODO: use this MAP while uploading to Huggingface
+# HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+#     "facebook/hubert-base-ls960": "https://huggingface.co/facebook/hubert-base-ls960/resolve/main/config.json",
+#     # See all Hubert models at https://huggingface.co/models?filter=hubert
+# }
+
+
+class MERTConfig(PretrainedConfig):
+    r"""
+    """
+    model_type = "mert_model"
+
+    def __init__(
+        self,
+        vocab_size=32,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout=0.1,
+        activation_dropout=0.1,
+        attention_dropout=0.1,
+        feat_proj_layer_norm=True,
+        feat_proj_dropout=0.0,
+        final_dropout=0.1,
+        layerdrop=0.1,
+        initializer_range=0.02,
+        layer_norm_eps=1e-5,
+        feat_extract_norm="group",
+        feat_extract_activation="gelu",
+        conv_dim=(512, 512, 512, 512, 512, 512, 512),
+        conv_stride=(5, 2, 2, 2, 2, 2, 2),
+        conv_kernel=(10, 3, 3, 3, 3, 2, 2),
+        conv_bias=False,
+        num_conv_pos_embeddings=128,
+        num_conv_pos_embedding_groups=16,
+        do_stable_layer_norm=False,
+        apply_spec_augment=True,
+        mask_time_prob=0.05,
+        mask_time_length=10,
+        mask_time_min_masks=2,
+        mask_feature_prob=0.0,
+        mask_feature_length=10,
+        mask_feature_min_masks=0,
+        ctc_loss_reduction="sum",
+        ctc_zero_infinity=False,
+        use_weighted_layer_sum=False,
+        classifier_proj_size=256,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        feature_extractor_cqt=False,
+        feature_extractor_cqt_bins=336,
+        deepnorm=False,
+        attention_relax=-1.0,
+        **kwargs
+    ):
+        super().__init__(**kwargs, pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id)
+        self.hidden_size = hidden_size
+        self.feat_extract_norm = feat_extract_norm
+        self.feat_extract_activation = feat_extract_activation
+        self.conv_dim = list(conv_dim)
+        self.conv_stride = list(conv_stride)
+        self.conv_kernel = list(conv_kernel)
+        self.conv_bias = conv_bias
+        self.num_conv_pos_embeddings = num_conv_pos_embeddings
+        self.num_conv_pos_embedding_groups = num_conv_pos_embedding_groups
+        self.num_feat_extract_layers = len(self.conv_dim)
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.num_attention_heads = num_attention_heads
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.feat_proj_layer_norm = feat_proj_layer_norm
+        self.feat_proj_dropout = feat_proj_dropout
+        self.final_dropout = final_dropout
+        self.layerdrop = layerdrop
+        self.layer_norm_eps = layer_norm_eps
+        self.initializer_range = initializer_range
+        self.vocab_size = vocab_size
+        self.do_stable_layer_norm = do_stable_layer_norm
+        self.use_weighted_layer_sum = use_weighted_layer_sum
+        self.classifier_proj_size = classifier_proj_size
+
+        if (
+            (len(self.conv_stride) != self.num_feat_extract_layers)
+            or (len(self.conv_kernel) != self.num_feat_extract_layers)
+            or (len(self.conv_dim) != self.num_feat_extract_layers)
+        ):
+            raise ValueError(
+                "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
+                " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
+                f" {len(self.conv_dim)}`, `len(config.conv_stride) = {len(self.conv_stride)}`,"
+                f" `len(config.conv_kernel) = {len(self.conv_kernel)}`."
+            )
+
+        # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
+        self.apply_spec_augment = apply_spec_augment
+        self.mask_time_prob = mask_time_prob
+        self.mask_time_length = mask_time_length
+        self.mask_time_min_masks = mask_time_min_masks
+        self.mask_feature_prob = mask_feature_prob
+        self.mask_feature_length = mask_feature_length
+        self.mask_feature_min_masks = mask_feature_min_masks
+
+        # ctc loss
+        self.ctc_loss_reduction = ctc_loss_reduction
+        self.ctc_zero_infinity = ctc_zero_infinity
+
+        # cqt feature extractor
+        self.feature_extractor_cqt = feature_extractor_cqt
+        self.feature_extractor_cqt_bins = feature_extractor_cqt_bins
+
+        # deepnorm: up-scale weighted residual conection + down-scale initial value transformer encoder
+        self.deepnorm = deepnorm
+
+        self.attention_relax = attention_relax
+
+    @property
+    def inputs_to_logits_ratio(self):
+        return functools.reduce(operator.mul, self.conv_stride, 1)

MERT-v1-95M/modeling_MERT.py

@@ -0,0 +1,409 @@
+"""
+MERT model definition.
+We largely adapt codes from:
+1. https://github.com/huggingface/transformers/blob/main/src/transformers/models/hubert/modeling_hubert.py
+2. https://github.com/facebookresearch/fairseq/blob/main/fairseq/models/wav2vec/wav2vec2.py
+"""
+
+from typing import Optional, Tuple, Union
+from transformers.modeling_outputs import BaseModelOutput
+import torch
+from torch import nn
+
+from transformers.models.hubert.modeling_hubert import (
+    HubertFeatureEncoder,
+    HubertModel,
+    HubertEncoderStableLayerNorm,
+    HubertEncoder,
+    HubertEncoderLayer,
+    HubertPositionalConvEmbedding,
+    HubertAttention,
+    HubertFeedForward,
+)
+
+try:
+    from nnAudio import features as nnAudioFeatures
+    NNAUDIO_INSTALLED=True
+except:
+    print("WARNING: feature_extractor_cqt requires the libray 'nnAudio'")
+    NNAUDIO_INSTALLED=False
+
+from .configuration_MERT import MERTConfig
+
+class MERTFeatureProjection(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.feat_proj_layer_norm = config.feat_proj_layer_norm
+        self.feature_extractor_cqt = config.feature_extractor_cqt
+
+        if self.feature_extractor_cqt:
+            # v3 concat features
+            self.feature_dimension = config.conv_dim[-1] + config.feature_extractor_cqt_bins
+            print(f"feature dimention: {self.feature_dimension}")
+        else:
+            self.feature_dimension = config.conv_dim[-1]
+        if self.feat_proj_layer_norm:
+            self.layer_norm = nn.LayerNorm(self.feature_dimension, eps=config.layer_norm_eps)
+        self.projection = nn.Linear(self.feature_dimension, config.hidden_size)
+        self.dropout = nn.Dropout(config.feat_proj_dropout)
+
+    def forward(self, hidden_states):
+        # non-projected hidden states are needed for quantization
+        if self.feat_proj_layer_norm:
+            hidden_states = self.layer_norm(hidden_states)
+        hidden_states = self.projection(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+class MERTModel(HubertModel):
+    # overwrite config class
+    config_class = MERTConfig
+    base_model_prefix = "mert_model"
+    def __init__(
+        self,
+        config: MERTConfig,
+    ) -> None:
+        """ 
+        initialize the with the grandparent method HubertPreTrainedModel.__init__()
+        and modify the HuBERTModel.__init__()
+        """
+        super(HubertModel, self).__init__(config)
+
+        self.config = config
+
+        self.feature_extractor = HubertFeatureEncoder(config)
+        self.feature_projection = MERTFeatureProjection(config) # replace Feature Projection for introcuing new feature
+
+        if self.config.feature_extractor_cqt:
+            assert NNAUDIO_INSTALLED, "ERROR: feature_extractor_cqt requires the libray 'nnAudio', try after `pip install nnAudio` "
+            print('initializing cqt extractor for MERT')            
+            self.feature_extractor_cqt = nnAudioFeatures.cqt.CQT(sr=self.config.sample_rate, hop_length=self.config.sample_rate//50, fmin=32.7, 
+                    fmax=None, n_bins=self.config.feature_extractor_cqt_bins, bins_per_octave=self.config.feature_extractor_cqt_bins//7, 
+                    filter_scale=1, norm=1, window='hann', center=True, 
+                    pad_mode='constant', trainable=False, 
+                    output_format='Magnitude', verbose=True)
+
+        if config.mask_time_prob > 0.0 or config.mask_feature_prob > 0.0:
+            self.masked_spec_embed = nn.Parameter(torch.FloatTensor(config.hidden_size).uniform_())
+
+        
+        if config.do_stable_layer_norm:
+            assert not config.deepnorm, "must use post-layer_norm with deepnorm"
+            self.encoder = HubertEncoderStableLayerNorm(config)
+        else:
+            if config.deepnorm:
+                self.encoder = HubertEncoder_extend(config)
+            else:
+                self.encoder = HubertEncoder(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+    
+    def forward(self, input_values: Optional[torch.Tensor], attention_mask: Optional[torch.Tensor] = None, mask_time_indices: Optional[torch.FloatTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None) -> Union[Tuple, BaseModelOutput]:
+        
+        # return super().forward(input_values, attention_mask, mask_time_indices, output_attentions, output_hidden_states, return_dict)
+        
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+
+        # add additional cqt features for transformer input
+        if self.config.feature_extractor_cqt:
+            features_cqt = self.feature_extractor_cqt(input_values).transpose(1, 2)
+            features_cqt = features_cqt[:,:extract_features.shape[1],:] # align shape
+            # # v2
+            # features_cqt = self.post_cqt_feature_proj(features_cqt)
+            # extract_features = self.feature_projection.layer_norm(extract_features) + self.feature_projection.layer_norm(features_cqt) #v2
+            # v3
+            extract_features = torch.cat([extract_features,features_cqt], 2)
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(extract_features.shape[1], attention_mask)
+
+        hidden_states = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0] # take last_hidden from encoder output
+
+        if not return_dict:
+            return (hidden_states,) + encoder_outputs[1:]
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+class HubertEncoder_extend(HubertEncoder):
+    def __init__(self, config):
+        # super().__init__()
+        # call nn module initialization
+        nn.Module.__init__(self)
+        # super(HubertEncoder_extend, self).__init__()
+
+        self.config = config
+        self.pos_conv_embed = HubertPositionalConvEmbedding(config)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout)
+
+        
+        self.layers = nn.ModuleList([HubertEncoderLayerExtend(config) for _ in range(config.num_hidden_layers)])
+
+        self.gradient_checkpointing = False
+
+        if config.deepnorm:
+            import math
+            init_scale = math.pow(8.0 * config.num_hidden_layers, 0.25)
+            for name, p in self.named_parameters():
+                if (
+                    "feed_forward.intermediate_dense" in name
+                    or "feed_forward.output_dense" in name
+                    or "out_proj" in name
+                    or "v_proj" in name
+                ):
+                    p.data.div_(init_scale)
+
+class HubertEncoderLayerExtend(HubertEncoderLayer):
+    def __init__(self, config):
+        nn.Module.__init__(self)
+        # super(HubertEncoderLayerExtend, self).__init__()
+        if config.attention_relax > 0 :
+            self.attention = HubertAttention_extend(
+                embed_dim=config.hidden_size,
+                num_heads=config.num_attention_heads,
+                dropout=config.attention_dropout,
+                is_decoder=False,
+                attention_relax=config.attention_relax,
+            )
+        else:    
+            self.attention = HubertAttention(
+                embed_dim=config.hidden_size,
+                num_heads=config.num_attention_heads,
+                dropout=config.attention_dropout,
+                is_decoder=False,
+            )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = HubertFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        if config.deepnorm:
+            import math
+            self.residual_alpha = math.pow(2.0 * config.num_hidden_layers, 0.25)
+        else:
+            self.residual_alpha = 1.0
+    
+    def residual_connection(self, x, residual):
+        '''
+        residual: input before f()
+        x: output of f(residual)
+        '''
+        return residual * self.residual_alpha + x
+
+    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
+        attn_residual = hidden_states
+        hidden_states, attn_weights, _ = self.attention(
+            hidden_states, attention_mask=attention_mask, output_attentions=output_attentions
+        )
+        hidden_states = self.dropout(hidden_states)
+
+        # hidden_states = attn_residual + hidden_states
+        hidden_states = self.residual_connection(hidden_states, attn_residual)
+
+        hidden_states = self.layer_norm(hidden_states)
+
+        # hidden_states = hidden_states + self.feed_forward(hidden_states)
+        ffn_residual = hidden_states
+        hidden_states = self.feed_forward(hidden_states)
+        hidden_states = self.residual_connection(hidden_states, ffn_residual)
+
+        hidden_states = self.final_layer_norm(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class HubertAttention_extend(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        is_decoder: bool = False,
+        bias: bool = True,
+        attention_relax: float = -1.0,
+    ):
+        super().__init__()
+        # nn.Module.__init__(self)
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.scaling = self.head_dim**-0.5
+        self.is_decoder = is_decoder
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if attention_relax > 0:
+            self.attention_relax = attention_relax
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        layer_head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+        is_cross_attention = key_value_states is not None
+
+        bsz, tgt_len, _ = hidden_states.size()
+
+        # get query proj
+        query_states = self.q_proj(hidden_states) * self.scaling
+        # get key, value proj
+        # `past_key_value[0].shape[2] == key_value_states.shape[1]`
+        # is checking that the `sequence_length` of the `past_key_value` is the same as
+        # the provided `key_value_states` to support prefix tuning
+        if (
+            is_cross_attention
+            and past_key_value is not None
+            and past_key_value[0].shape[2] == key_value_states.shape[1]
+        ):
+            # reuse k,v, cross_attentions
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+        elif is_cross_attention:
+            # cross_attentions
+            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
+            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
+        elif past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        else:
+            # self_attention
+            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+        if self.is_decoder:
+            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+            # Further calls to cross_attention layer can then reuse all cross-attention
+            # key/value_states (first "if" case)
+            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+            # all previous decoder key/value_states. Further calls to uni-directional self-attention
+            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+            # if encoder bi-directional self-attention `past_key_value` is always `None`
+            past_key_value = (key_states, value_states)
+
+        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
+        key_states = key_states.view(*proj_shape)
+        value_states = value_states.view(*proj_shape)
+
+        src_len = key_states.size(1)
+        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if self.attention_relax > 0:
+            # => (bsz, self.num_heads, tgt_len, src_len)
+            # attn_weights_relax = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)/self.attention_relax
+            # => (bsz*self.num_heads, tgt_len, src_len)
+            attn_weights_relax = attn_weights / self.attention_relax
+
+            # => (bsz* self.num_heads, tgt_len, 1)
+            attn_max_relax = torch.max(attn_weights_relax, dim=-1, keepdim=False).unsqueeze(2)
+            attn_weights = (attn_weights_relax - attn_max_relax) * self.attention_relax
+
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+        if layer_head_mask is not None:
+            if layer_head_mask.size() != (self.num_heads,):
+                raise ValueError(
+                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
+                    f" {layer_head_mask.size()}"
+                )
+            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if output_attentions:
+            # this operation is a bit awkward, but it's required to
+            # make sure that attn_weights keeps its gradient.
+            # In order to do so, attn_weights have to be reshaped
+            # twice and have to be reused in the following
+            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
+        else:
+            attn_weights_reshaped = None
+
+        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.bmm(attn_probs, value_states)
+
+        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+        attn_output = attn_output.transpose(1, 2)
+
+        # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
+        # partitioned aross GPUs when using tensor-parallelism.
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights_reshaped, past_key_value

MERT-v1-95M/preprocessor_config.json

@@ -0,0 +1,9 @@
+{
+  "do_normalize": false,
+  "feature_extractor_type": "Wav2Vec2FeatureExtractor",
+  "feature_size": 1,
+  "padding_side": "right",
+  "padding_value": 0,
+  "return_attention_mask": true,
+  "sampling_rate": 24000
+}

MERT-v1-95M/pytorch_model.bin

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

Prediction_Head/MTGGenre_head.py

@@ -0,0 +1,48 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+class MLPProberBase(nn.Module):
+    def __init__(self, d=768, layer='all', num_outputs=87):
+        super().__init__()
+        
+        self.hidden_layer_sizes = [512, ] # eval(self.cfg.hidden_layer_sizes)
+        
+        self.num_layers = len(self.hidden_layer_sizes)
+
+        self.layer = layer
+
+        for i, ld in enumerate(self.hidden_layer_sizes):
+            setattr(self, f"hidden_{i}", nn.Linear(d, ld))
+            d = ld
+        self.output = nn.Linear(d, num_outputs)
+
+        self.n_tranformer_layer = 12
+        
+        self.init_aggregator()
+
+
+    def init_aggregator(self):
+        """Initialize the aggregator for weighted sum over different layers of features
+        """
+        if self.layer == "all":
+            # use learned weights to aggregate features
+            self.aggregator = nn.Parameter(torch.randn((1, self.n_tranformer_layer, 1)))
+
+
+    def forward(self, x):
+        """
+        x: (B, L, T, H)
+        T=#chunks, can be 1 or several chunks
+        """
+        
+        if self.layer == "all":
+            weights = F.softmax(self.aggregator, dim=1)
+            x = (x * weights).sum(dim=1)
+
+        for i in range(self.num_layers):
+            x = getattr(self, f"hidden_{i}")(x)
+            # x = self.dropout(x)
+            x = F.relu(x)
+        output = self.output(x)
+        return output

Prediction_Head/MTGGenre_id2class.json

@@ -0,0 +1 @@
+{"0": "genre---rock", "1": "genre---pop", "2": "genre---classical", "3": "genre---popfolk", "4": "genre---disco", "5": "genre---funk", "6": "genre---rnb", "7": "genre---ambient", "8": "genre---chillout", "9": "genre---downtempo", "10": "genre---easylistening", "11": "genre---electronic", "12": "genre---lounge", "13": "genre---triphop", "14": "genre---breakbeat", "15": "genre---techno", "16": "genre---newage", "17": "genre---jazz", "18": "genre---metal", "19": "genre---industrial", "20": "genre---instrumentalrock", "21": "genre---minimal", "22": "genre---alternative", "23": "genre---experimental", "24": "genre---drumnbass", "25": "genre---soul", "26": "genre---fusion", "27": "genre---soundtrack", "28": "genre---electropop", "29": "genre---world", "30": "genre---ethno", "31": "genre---trance", "32": "genre---orchestral", "33": "genre---grunge", "34": "genre---chanson", "35": "genre---worldfusion", "36": "genre---hiphop", "37": "genre---groove", "38": "genre---instrumentalpop", "39": "genre---blues", "40": "genre---reggae", "41": "genre---dance", "42": "genre---club", "43": "genre---punkrock", "44": "genre---folk", "45": "genre---synthpop", "46": "genre---poprock", "47": "genre---choir", "48": "genre---symphonic", "49": "genre---indie", "50": "genre---progressive", "51": "genre---acidjazz", "52": "genre---contemporary", "53": "genre---newwave", "54": "genre---dub", "55": "genre---rocknroll", "56": "genre---hard", "57": "genre---hardrock", "58": "genre---house", "59": "genre---atmospheric", "60": "genre---psychedelic", "61": "genre---improvisation", "62": "genre---country", "63": "genre---electronica", "64": "genre---rap", "65": "genre---60s", "66": "genre---70s", "67": "genre---darkambient", "68": "genre---idm", "69": "genre---latin", "70": "genre---postrock", "71": "genre---bossanova", "72": "genre---singersongwriter", "73": "genre---darkwave", "74": "genre---swing", "75": "genre---medieval", "76": "genre---celtic", "77": "genre---eurodance", "78": "genre---classicrock", "79": "genre---dubstep", "80": "genre---bluesrock", "81": "genre---edm", "82": "genre---deephouse", "83": "genre---jazzfusion", "84": "genre---alternativerock", "85": "genre---80s", "86": "genre---90s"}

Prediction_Head/best-layer-MERT-v1-95M.zip

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

Prediction_Head/best-layer-MERT-v1-95M/EMO.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/EMO.id2class.json

@@ -0,0 +1 @@
+{"0": "arousal", "1": "valence"}

Prediction_Head/best-layer-MERT-v1-95M/GS.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/GS.id2class.json

@@ -0,0 +1 @@
+{"0": "C major", "1": "Db major", "2": "D major", "3": "Eb major", "4": "E major", "5": "F major", "6": "Gb major", "7": "G major", "8": "Ab major", "9": "A major", "10": "Bb major", "11": "B major", "12": "C minor", "13": "Db minor", "14": "D minor", "15": "Eb minor", "16": "E minor", "17": "F minor", "18": "Gb minor", "19": "G minor", "20": "Ab minor", "21": "A minor", "22": "Bb minor", "23": "B minor"}

Prediction_Head/best-layer-MERT-v1-95M/GTZAN.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/MTGGenre.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/MTGGenre.id2class.json

@@ -0,0 +1 @@
+{"0": "genre---rock", "1": "genre---pop", "2": "genre---classical", "3": "genre---popfolk", "4": "genre---disco", "5": "genre---funk", "6": "genre---rnb", "7": "genre---ambient", "8": "genre---chillout", "9": "genre---downtempo", "10": "genre---easylistening", "11": "genre---electronic", "12": "genre---lounge", "13": "genre---triphop", "14": "genre---breakbeat", "15": "genre---techno", "16": "genre---newage", "17": "genre---jazz", "18": "genre---metal", "19": "genre---industrial", "20": "genre---instrumentalrock", "21": "genre---minimal", "22": "genre---alternative", "23": "genre---experimental", "24": "genre---drumnbass", "25": "genre---soul", "26": "genre---fusion", "27": "genre---soundtrack", "28": "genre---electropop", "29": "genre---world", "30": "genre---ethno", "31": "genre---trance", "32": "genre---orchestral", "33": "genre---grunge", "34": "genre---chanson", "35": "genre---worldfusion", "36": "genre---hiphop", "37": "genre---groove", "38": "genre---instrumentalpop", "39": "genre---blues", "40": "genre---reggae", "41": "genre---dance", "42": "genre---club", "43": "genre---punkrock", "44": "genre---folk", "45": "genre---synthpop", "46": "genre---poprock", "47": "genre---choir", "48": "genre---symphonic", "49": "genre---indie", "50": "genre---progressive", "51": "genre---acidjazz", "52": "genre---contemporary", "53": "genre---newwave", "54": "genre---dub", "55": "genre---rocknroll", "56": "genre---hard", "57": "genre---hardrock", "58": "genre---house", "59": "genre---atmospheric", "60": "genre---psychedelic", "61": "genre---improvisation", "62": "genre---country", "63": "genre---electronica", "64": "genre---rap", "65": "genre---60s", "66": "genre---70s", "67": "genre---darkambient", "68": "genre---idm", "69": "genre---latin", "70": "genre---postrock", "71": "genre---bossanova", "72": "genre---singersongwriter", "73": "genre---darkwave", "74": "genre---swing", "75": "genre---medieval", "76": "genre---celtic", "77": "genre---eurodance", "78": "genre---classicrock", "79": "genre---dubstep", "80": "genre---bluesrock", "81": "genre---edm", "82": "genre---deephouse", "83": "genre---jazzfusion", "84": "genre---alternativerock", "85": "genre---80s", "86": "genre---90s"}

Prediction_Head/best-layer-MERT-v1-95M/MTGInstrument.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/MTGInstrument.id2class.json

@@ -0,0 +1 @@
+{"0": "instrument---voice", "1": "instrument---synthesizer", "2": "instrument---piano", "3": "instrument---strings", "4": "instrument---beat", "5": "instrument---guitar", "6": "instrument---cello", "7": "instrument---keyboard", "8": "instrument---trombone", "9": "instrument---clarinet", "10": "instrument---doublebass", "11": "instrument---horn", "12": "instrument---trumpet", "13": "instrument---violin", "14": "instrument---accordion", "15": "instrument---bass", "16": "instrument---computer", "17": "instrument---drummachine", "18": "instrument---drums", "19": "instrument---electricguitar", "20": "instrument---sampler", "21": "instrument---acousticguitar", "22": "instrument---harmonica", "23": "instrument---flute", "24": "instrument---pipeorgan", "25": "instrument---harp", "26": "instrument---electricpiano", "27": "instrument---oboe", "28": "instrument---saxophone", "29": "instrument---percussion", "30": "instrument---acousticbassguitar", "31": "instrument---orchestra", "32": "instrument---bongo", "33": "instrument---brass", "34": "instrument---viola", "35": "instrument---rhodes", "36": "instrument---organ", "37": "instrument---classicalguitar", "38": "instrument---bell", "39": "instrument---pad"}

Prediction_Head/best-layer-MERT-v1-95M/MTGMood.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/MTGMood.id2class.json

@@ -0,0 +1 @@
+{"0": "mood/theme---background", "1": "mood/theme---film", "2": "mood/theme---melancholic", "3": "mood/theme---melodic", "4": "mood/theme---children", "5": "mood/theme---relaxing", "6": "mood/theme---documentary", "7": "mood/theme---emotional", "8": "mood/theme---space", "9": "mood/theme---love", "10": "mood/theme---drama", "11": "mood/theme---adventure", "12": "mood/theme---energetic", "13": "mood/theme---heavy", "14": "mood/theme---dark", "15": "mood/theme---calm", "16": "mood/theme---action", "17": "mood/theme---dramatic", "18": "mood/theme---epic", "19": "mood/theme---powerful", "20": "mood/theme---upbeat", "21": "mood/theme---slow", "22": "mood/theme---inspiring", "23": "mood/theme---soft", "24": "mood/theme---meditative", "25": "mood/theme---fun", "26": "mood/theme---happy", "27": "mood/theme---positive", "28": "mood/theme---romantic", "29": "mood/theme---sad", "30": "mood/theme---hopeful", "31": "mood/theme---motivational", "32": "mood/theme---deep", "33": "mood/theme---uplifting", "34": "mood/theme---ballad", "35": "mood/theme---soundscape", "36": "mood/theme---dream", "37": "mood/theme---movie", "38": "mood/theme---fast", "39": "mood/theme---nature", "40": "mood/theme---cool", "41": "mood/theme---corporate", "42": "mood/theme---travel", "43": "mood/theme---funny", "44": "mood/theme---sport", "45": "mood/theme---commercial", "46": "mood/theme---advertising", "47": "mood/theme---holiday", "48": "mood/theme---christmas", "49": "mood/theme---sexy", "50": "mood/theme---game", "51": "mood/theme---groovy", "52": "mood/theme---retro", "53": "mood/theme---summer", "54": "mood/theme---party", "55": "mood/theme---trailer"}

Prediction_Head/best-layer-MERT-v1-95M/MTGTop50.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/MTGTop50.id2class.json

@@ -0,0 +1 @@
+{"0": "genre---rock", "1": "genre---pop", "2": "genre---classical", "3": "instrument---voice", "4": "genre---popfolk", "5": "genre---funk", "6": "genre---ambient", "7": "genre---chillout", "8": "genre---downtempo", "9": "genre---easylistening", "10": "genre---electronic", "11": "genre---lounge", "12": "instrument---synthesizer", "13": "genre---triphop", "14": "genre---techno", "15": "genre---newage", "16": "genre---jazz", "17": "genre---metal", "18": "instrument---piano", "19": "genre---alternative", "20": "genre---experimental", "21": "genre---soundtrack", "22": "mood/theme---film", "23": "genre---world", "24": "instrument---strings", "25": "genre---trance", "26": "genre---orchestral", "27": "instrument---guitar", "28": "genre---hiphop", "29": "genre---instrumentalpop", "30": "mood/theme---relaxing", "31": "genre---reggae", "32": "mood/theme---emotional", "33": "instrument---keyboard", "34": "instrument---violin", "35": "genre---dance", "36": "instrument---bass", "37": "instrument---computer", "38": "instrument---drummachine", "39": "instrument---drums", "40": "instrument---electricguitar", "41": "genre---folk", "42": "instrument---acousticguitar", "43": "genre---poprock", "44": "genre---indie", "45": "mood/theme---energetic", "46": "mood/theme---happy", "47": "instrument---electricpiano", "48": "genre---house", "49": "genre---atmospheric"}

Prediction_Head/best-layer-MERT-v1-95M/NSynthI.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/NSynthI.id2class.json

@@ -0,0 +1 @@
+{"0": "bass", "1": "brass", "2": "flute", "3": "guitar", "4": "keyboard", "5": "mallet", "6": "organ", "7": "reed", "8": "string", "9": "synth_lead", "10": "vocal"}

Prediction_Head/best-layer-MERT-v1-95M/NSynthP.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/NSynthP.id2class.json

@@ -0,0 +1 @@
+{"0": 1, "1": 2, "2": 3, "3": 4, "4": 5, "5": 6, "6": 7, "7": 8, "8": 9, "9": 10, "10": 11, "11": 12, "12": 13, "13": 14, "14": 15, "15": 16, "16": 17, "17": 18, "18": 19, "19": 20, "20": 21, "21": 22, "22": 23, "23": 24, "24": 25, "25": 26, "26": 27, "27": 28, "28": 29, "29": 30, "30": 31, "31": 32, "32": 33, "33": 34, "34": 35, "35": 36, "36": 37, "37": 38, "38": 39, "39": 40, "40": 41, "41": 42, "42": 43, "43": 44, "44": 45, "45": 46, "46": 47, "47": 48, "48": 49, "49": 50, "50": 51, "51": 52, "52": 53, "53": 54, "54": 55, "55": 56, "56": 57, "57": 58, "58": 59, "59": 60, "60": 61, "61": 62, "62": 63, "63": 64, "64": 65, "65": 66, "66": 67, "67": 68, "68": 69, "69": 70, "70": 71, "71": 72, "72": 73, "73": 74, "74": 75, "75": 76, "76": 77, "77": 78, "78": 79, "79": 80, "80": 81, "81": 82, "82": 83, "83": 84, "84": 85, "85": 86, "86": 87, "87": 88, "88": 89, "89": 90, "90": 91, "91": 92, "92": 93, "93": 94, "94": 95, "95": 96, "96": 97, "97": 98, "98": 99, "99": 100, "100": 101, "101": 102, "102": 103, "103": 104, "104": 105, "105": 106, "106": 107, "107": 108, "108": 109, "109": 110, "110": 111, "111": 112, "112": 113, "113": 114, "114": 115, "115": 116, "116": 117, "117": 118, "118": 119, "119": 120, "120": 121, "121": 122, "122": 123, "123": 124, "124": 125, "125": 126, "126": 127, "127": 128}

Prediction_Head/best-layer-MERT-v1-95M/VocalSetS.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/VocalSetS.id2class.json

@@ -0,0 +1 @@
+{"0": "f1", "1": "f2", "2": "f3", "3": "f4", "4": "f5", "5": "f6", "6": "f7", "7": "f8", "8": "f9", "9": "m1", "10": "m2", "11": "m3", "12": "m4", "13": "m5", "14": "m6", "15": "m7", "16": "m8", "17": "m9", "18": "m10", "19": "m11"}

Prediction_Head/best-layer-MERT-v1-95M/VocalSetT.ckpt

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

Prediction_Head/best-layer-MERT-v1-95M/VocalSetT.id2class.json

@@ -0,0 +1 @@
+{"0": "belt", "1": "breathy", "2": "inhaled", "3": "lip_trill", "4": "spoken", "5": "straight", "6": "trill", "7": "trillo", "8": "vibrato", "9": "vocal_fry"}

Prediction_Head/best_MTGGenre.ckpt

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

README.md

@@ -1,10 +1,10 @@
 ---
-title: Music Descriptor
-emoji: 📚
-colorFrom: yellow
-colorTo: green
+title: MusicTagging
+emoji: 💻
+colorFrom: gray
+colorTo: red
 sdk: gradio
-sdk_version: 3.32.0
+sdk_version: 3.29.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -0,0 +1,207 @@
+import gradio as gr
+#
+from transformers import Wav2Vec2FeatureExtractor
+from transformers import AutoModel
+import torch
+from torch import nn
+import torchaudio
+import torchaudio.transforms as T
+import logging
+
+import json
+import os
+
+import importlib 
+modeling_MERT = importlib.import_module("MERT-v1-95M.modeling_MERT")
+
+from Prediction_Head.MTGGenre_head import MLPProberBase 
+# input cr: https://huggingface.co/spaces/thealphhamerc/audio-to-text/blob/main/app.py
+
+
+logger = logging.getLogger("MERT-v1-95M-app")
+logger.setLevel(logging.INFO)
+ch = logging.StreamHandler()
+ch.setLevel(logging.INFO)
+formatter = logging.Formatter(
+    "%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d %H:%M:%S")
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+
+
+inputs = [
+    gr.components.Audio(type="filepath", label="Add music audio file"), 
+    gr.inputs.Audio(source="microphone", type="filepath"),
+]
+live_inputs = [
+    gr.Audio(source="microphone",streaming=True, type="filepath"),
+]
+# outputs = [gr.components.Textbox()]
+# outputs = [gr.components.Textbox(), transcription_df]
+title = "One Model for All Music Understanding Tasks"
+description = "An example of using the [MERT-v1-95M](https://huggingface.co/m-a-p/MERT-v1-95M) model as backbone to conduct multiple music understanding tasks with the universal represenation."
+article = "The tasks include EMO, GS, MTGInstrument, MTGGenre, MTGTop50, MTGMood, NSynthI, NSynthP, VocalSetS, VocalSetT. \n\n More models can be referred at the [map organization page](https://huggingface.co/m-a-p)."
+audio_examples = [
+    # ["input/example-1.wav"],
+    # ["input/example-2.wav"],
+]
+
+# Load the model and the corresponding preprocessor config
+# model = AutoModel.from_pretrained("m-a-p/MERT-v0-public", trust_remote_code=True)
+# processor = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v0-public",trust_remote_code=True)
+model = modeling_MERT.MERTModel.from_pretrained("./MERT-v1-95M")
+processor = Wav2Vec2FeatureExtractor.from_pretrained("./MERT-v1-95M")
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+MERT_BEST_LAYER_IDX = {
+    'EMO': 5,
+    'GS': 8,
+    'GTZAN': 7,
+    'MTGGenre': 7,
+    'MTGInstrument': 'all',
+    'MTGMood': 6,
+    'MTGTop50': 6,
+    'MTT': 'all',
+    'NSynthI': 6,
+    'NSynthP': 1,
+    'VocalSetS': 2,
+    'VocalSetT': 9,
+} 
+
+MERT_BEST_LAYER_IDX = {
+    'EMO': 5,
+    'GS': 8,
+    'GTZAN': 7,
+    'MTGGenre': 7,
+    'MTGInstrument': 'all',
+    'MTGMood': 6,
+    'MTGTop50': 6,
+    'MTT': 'all',
+    'NSynthI': 6,
+    'NSynthP': 1,
+    'VocalSetS': 2,
+    'VocalSetT': 9,
+} 
+CLASSIFIERS = {
+
+}
+
+ID2CLASS = {
+
+}
+
+TASKS = ['EMO','GS', 'MTGInstrument', 'MTGGenre', 'MTGTop50', 'MTGMood', 'NSynthI', 'NSynthP', 'VocalSetS', 'VocalSetT']
+head_dir = '/home/chenghua/nanshen/Yizhi/MERT_Universal/Prediction_Head/best-layer-MERT-v1-95M'
+for task in TASKS:
+    print('loading', task)
+    with open(os.path.join(head_dir,f'{task}.id2class.json'), 'r') as f:
+        ID2CLASS[task]=json.load(f)
+    num_class = len(ID2CLASS[task].keys())
+    CLASSIFIERS[task] = MLPProberBase(d=768, layer=MERT_BEST_LAYER_IDX[task], num_outputs=num_class)
+    CLASSIFIERS[task].load_state_dict(torch.load(f'/home/chenghua/nanshen/Yizhi/MERT_Universal/Prediction_Head/best-layer-MERT-v1-95M/{task}.ckpt')['state_dict'])
+    CLASSIFIERS[task].to(device)
+
+model.to(device)
+
+def model_infernce(inputs):
+    waveform, sample_rate = torchaudio.load(inputs)
+
+    resample_rate = processor.sampling_rate
+
+    # make sure the sample_rate aligned
+    if resample_rate != sample_rate:
+        print(f'setting rate from {sample_rate} to {resample_rate}')
+        resampler = T.Resample(sample_rate, resample_rate)
+        waveform = resampler(waveform)
+    
+    waveform = waveform.view(-1,) # make it (n_sample, )
+    model_inputs = processor(waveform, sampling_rate=resample_rate, return_tensors="pt")
+    model_inputs.to(device)
+    with torch.no_grad():
+        model_outputs = model(**model_inputs, output_hidden_states=True)
+
+    # take a look at the output shape, there are 13 layers of representation
+    # each layer performs differently in different downstream tasks, you should choose empirically
+    all_layer_hidden_states = torch.stack(model_outputs.hidden_states).squeeze()[1:,:,:].unsqueeze(0)
+    print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim]
+    all_layer_hidden_states = all_layer_hidden_states.mean(dim=2)
+
+    task_output_texts = ""
+    for task in TASKS:
+        num_class = len(ID2CLASS[task].keys())
+        if MERT_BEST_LAYER_IDX[task] == 'all':
+            logits = CLASSIFIERS[task](all_layer_hidden_states) # [1, 87]
+        else:
+            logits = CLASSIFIERS[task](all_layer_hidden_states[:, MERT_BEST_LAYER_IDX[task]])
+        print(f'task {task} logits:', logits.shape, 'num class:', num_class)
+        
+        sorted_idx = torch.argsort(logits, dim = -1, descending=True)[0] # batch =1 
+        sorted_prob,_ = torch.sort(nn.functional.softmax(logits[0], dim=-1), dim=-1, descending=True)
+        # print(sorted_prob)
+        # print(sorted_prob.shape)
+        
+        top_n_show = 3 if num_class >= 3 else num_class
+        task_output_texts = task_output_texts + f"TASK {task} output:\n" + "\n".join([str(ID2CLASS[task][str(sorted_idx[idx].item())])+f', probability: {sorted_prob[idx].item():.2%}' for idx in range(top_n_show)]) + '\n'
+        task_output_texts = task_output_texts + '----------------------\n'
+        # output_texts = "\n".join([id2cls[str(idx.item())].replace('genre---', '') for idx in sorted_idx[:5]])
+    # logger.warning(all_layer_hidden_states.shape)
+    
+    # return f"device {device}, sample reprensentation:  {str(all_layer_hidden_states[12, 0, :10])}"
+    # return f"device: {device}\n" + output_texts
+    return task_output_texts
+    
+def convert_audio(inputs, microphone):
+    if (microphone is not None):
+        inputs = microphone
+
+    text = model_infernce(inputs)
+
+    return text
+    
+def live_convert_audio(microphone):
+    if (microphone is not None):
+        inputs = microphone
+    
+    text = model_infernce(inputs)
+
+    return text
+
+audio_chunked = gr.Interface(
+    fn=convert_audio,
+    inputs=inputs,
+    outputs=[gr.components.Textbox()],
+    allow_flagging="never",
+    title=title,
+    description=description,
+    article=article,
+    examples=audio_examples,
+)
+
+live_audio_chunked = gr.Interface(
+    fn=live_convert_audio,
+    inputs=live_inputs,
+    outputs=[gr.components.Textbox()],
+    allow_flagging="never",
+    title=title,
+    description=description,
+    article=article,
+    # examples=audio_examples,
+    live=True,
+)
+
+
+demo = gr.Blocks()
+with demo:
+    gr.TabbedInterface(
+        [
+            audio_chunked,
+            live_audio_chunked,
+        ], 
+        [
+            "Audio File or Recording",
+            "Live Streaming Music"
+        ]
+    )
+demo.queue(concurrency_count=1, max_size=5)
+demo.launch(show_api=False)

requirements.txt

@@ -0,0 +1,88 @@
+aiofiles==23.1.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==5.0.0
+anyio==3.6.2
+async-timeout==4.0.2
+attrs==23.1.0
+certifi==2023.5.7
+charset-normalizer==3.1.0
+click==8.1.3
+cmake==3.26.3
+contourpy==1.0.7
+cycler==0.11.0
+fastapi==0.95.2
+ffmpy==0.3.0
+filelock==3.12.0
+fonttools==4.39.4
+frozenlist==1.3.3
+fsspec==2023.5.0
+gradio==3.31.0
+gradio_client==0.2.5
+h11==0.14.0
+httpcore==0.17.1
+httpx==0.24.0
+huggingface-hub==0.14.1
+idna==3.4
+Jinja2==3.1.2
+jsonschema==4.17.3
+kiwisolver==1.4.4
+linkify-it-py==2.0.2
+lit==16.0.5
+markdown-it-py==2.2.0
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.1
+nnAudio==0.3.2
+numpy==1.24.3
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+orjson==3.8.12
+packaging==23.1
+pandas==2.0.1
+Pillow==9.5.0
+pydantic==1.10.7
+pydub==0.25.1
+Pygments==2.15.1
+pyparsing==3.0.9
+pyrsistent==0.19.3
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3
+PyYAML==6.0
+regex==2023.5.5
+requests==2.30.0
+scipy==1.10.1
+semantic-version==2.10.0
+six==1.16.0
+sniffio==1.3.0
+starlette==0.27.0
+sympy==1.12
+tokenizers==0.13.3
+toolz==0.12.0
+torch==2.0.1
+torchaudio==2.0.2
+torchvision==0.15.2
+tqdm==4.65.0
+transformers==4.29.2
+triton==2.0.0
+typing_extensions==4.5.0
+tzdata==2023.3
+uc-micro-py==1.0.2
+urllib3==2.0.2
+uvicorn==0.22.0
+websockets==11.0.3
+yarl==1.9.2