add genre prediction head

Prediction_Head/MTGGenre_head.py

@@ -0,0 +1,21 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+class MLPProberBase(nn.Module):
+    def __init__(self, d=768, num_outputs=87):
+        super().__init__()
+        self.hidden_layer_sizes = [512, ] # eval(self.cfg.hidden_layer_sizes)
+        self.num_layers = len(self.hidden_layer_sizes)
+        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)
+    
+    def forward(self, x):
+        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_MTGGenre.ckpt

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

app.py

@@ -8,9 +8,12 @@ import torchaudio
 import torchaudio.transforms as T
 import logging
 
+import json
+
 import importlib 
 modeling_MERT = importlib.import_module("MERT-v0-public.modeling_MERT")
 
+from Prediction_Head.MTGGenre_head import MLPProberBase 
 # input cr: https://huggingface.co/spaces/thealphhamerc/audio-to-text/blob/main/app.py
 
 
@@ -34,7 +37,7 @@ live_inputs = [
 ]
 # outputs = [gr.components.Textbox()]
 # outputs = [gr.components.Textbox(), transcription_df]
-title = "Output the tags of a (music) audio"
+title = "Predict the top 5 possible genres of Music"
 description = "An example of using MERT-95M-public to conduct music tagging."
 article = ""
 audio_examples = [
@@ -48,9 +51,17 @@ audio_examples = [
 model = modeling_MERT.MERTModel.from_pretrained("./MERT-v0-public")
 processor = Wav2Vec2FeatureExtractor.from_pretrained("./MERT-v0-public")
 
+MERT_LAYER_IDX = 7
+MTGGenre_classifier = MLPProberBase()
+MTGGenre_classifier.load_state_dict(torch.load('Prediction_Head/best_MTGGenre.ckpt')['state_dict'])
+
+with open('Prediction_Head/MTGGenre_id2class.json', 'r') as f:
+   id2cls=json.load(f)
+
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model.to(device)
+MTGGenre_classifier.to(device)
 
 def convert_audio(inputs, microphone):
     if (microphone is not None):
@@ -75,10 +86,17 @@ def convert_audio(inputs, microphone):
     # 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()
-    # print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim]
+    print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim]
+
+    logits = MTGGenre_classifier(torch.mean(all_layer_hidden_states[MERT_LAYER_IDX], dim=0)) # [1, 87]
+    print(logits.shape)
+    sorted_idx = torch.argsort(logits, dim = -1, descending=True)
+    
+    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}\n sample reprensentation:  {str(all_layer_hidden_states[12, 0, :10])}"
+    # return f"device {device}, sample reprensentation:  {str(all_layer_hidden_states[12, 0, :10])}"
+    return f"device: {device}\n" + output_texts
 
 def live_convert_audio(microphone):
     if (microphone is not None):
@@ -103,10 +121,17 @@ def live_convert_audio(microphone):
     # 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()
-    # print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim]
+    print(all_layer_hidden_states.shape) # [13 layer, Time steps, 768 feature_dim]
+
+    logits = MTGGenre_classifier(torch.mean(all_layer_hidden_states[MERT_LAYER_IDX], dim=0)) # [1, 87]
+    print(logits.shape)
+    sorted_idx = torch.argsort(logits, dim = -1, descending=True)
+    
+    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}, sample reprensentation:  {str(all_layer_hidden_states[12, 0, :10])}"
+    return f"device: {device}\n" + output_texts
 
 
 audio_chunked = gr.Interface(