Upload 2 files

model.py

@@ -0,0 +1,185 @@
+from torch.ao.nn.quantized import Sigmoid
+from transformers import BartModel
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+from peft import get_peft_model, LoraConfig
+from huggingface_hub import PyTorchModelHubMixin
+from transformers import BartConfig
+
+
+class MLP(nn.Module):
+    def __init__(self, layer_sizes=[64, 64, 64, 1], arl=False, dropout=0.0):
+        super().__init__()
+        self.arl = arl
+        self.attention = nn.Sequential(
+            nn.Linear(layer_sizes[0], layer_sizes[0]),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(layer_sizes[0], layer_sizes[0])
+        )
+
+        self.layer_sizes = layer_sizes
+        if len(layer_sizes) < 2:
+            raise ValueError()
+        self.layers = nn.ModuleList()
+        self.act = nn.LeakyReLU(negative_slope=0.01, inplace=True)
+        self.dropout = nn.Dropout(dropout)
+        for i in range(len(layer_sizes) - 1):
+            self.layers.append(nn.Linear(layer_sizes[i], layer_sizes[i + 1]))
+
+    def forward(self, x):
+        if self.arl:
+            x = x * self.attention(x)
+        for layer in self.layers[:-1]:
+            x = self.dropout(self.act(layer(x)))
+        x = self.layers[-1](x)
+        return x
+
+
+class BART(nn.Module):
+    def __init__(self, bartconfig, class_num=100):
+        super().__init__()
+        d_model = bartconfig.d_model
+        self.decoder_emb = nn.Embedding(class_num, d_model)
+        self.bart = BartModel(bartconfig)
+
+    def forward(self, x_encoder, x_decoder, attn_mask_encoder=None, attn_mask_decoder=None):
+        emb_encoder = x_encoder
+        emb_decoder = self.decoder_emb(x_decoder)
+        y = self.bart(inputs_embeds=emb_encoder, decoder_inputs_embeds=emb_decoder,
+                      attention_mask=attn_mask_encoder, decoder_attention_mask=attn_mask_decoder,
+                      output_hidden_states=False)
+        y = y.last_hidden_state
+        return y
+
+    def encode(self, x_encoder, attn_mask_encoder=None):
+        emb_encoder = x_encoder
+        y = self.bart.encoder(inputs_embeds=emb_encoder, attention_mask=attn_mask_encoder, output_hidden_states=False)
+        y = y.last_hidden_state
+        return y
+
+
+class ML_BART(nn.Module):
+    def __init__(self, bartconfig, class_num=[180, 256], pretrain=False, music_dim=512):
+        super().__init__()
+        d_model = bartconfig.d_model
+
+        self.decoder_emb = nn.ModuleList([
+            nn.Embedding(class_num[0] + 1, d_model // 4),
+            nn.Embedding(class_num[1] + 1, d_model // 4)
+        ])
+        self.decoder = MLP([music_dim, d_model // 2])
+
+        self.bart = BartModel(bartconfig)
+        self.pretrain = pretrain
+
+        self.encoder = MLP([music_dim, d_model])
+        self.lora_config = LoraConfig(
+            r=4,
+            lora_alpha=16,
+            lora_dropout=0.1
+        )
+
+    def forward(self, x_encoder, x_decoder, attn_mask_encoder=None, attn_mask_decoder=None):
+        # emb_encoder = x_encoder
+        emb_encoder = self.encoder(x_encoder)
+
+        if self.pretrain:
+            # emb_decoder = x_decoder
+            emb_decoder = self.encoder(x_decoder)
+        else:
+            emb_decoder = torch.concatenate(
+                [self.decoder_emb[0](x_decoder[..., 0]), self.decoder_emb[1](x_decoder[..., 1]),
+                 self.decoder(x_encoder)], dim=-1)
+
+        y = self.bart(inputs_embeds=emb_encoder, decoder_inputs_embeds=emb_decoder,
+                      attention_mask=attn_mask_encoder, decoder_attention_mask=attn_mask_decoder,
+                      output_hidden_states=False)
+        y = y.last_hidden_state
+        return y
+
+    def encode(self, x_encoder, attn_mask_encoder=None):
+        # emb_encoder = x_encoder
+        emb_encoder = self.encoder(x_encoder)
+
+        y = self.bart.encoder(inputs_embeds=emb_encoder, attention_mask=attn_mask_encoder, output_hidden_states=False)
+        y = y.last_hidden_state
+        return y
+
+    def reset_decoder(self):
+        for name, param in self.bart.decoder.named_parameters():
+            if param.dim() >= 2:
+                init.xavier_uniform_(param)
+            elif param.dim() == 1:
+                init.zeros_(param)
+
+
+class ML_Classifier(nn.Module):
+    def __init__(self, hidden_dim=512, class_num=[180, 256]):
+        super().__init__()
+        self.classifier = nn.ModuleList([
+            MLP([hidden_dim, hidden_dim, class_num[0] + 1]),
+            MLP([hidden_dim, hidden_dim, class_num[1] + 1])
+        ])
+
+    def forward(self, x):
+        h = self.classifier[0](x)
+        v = self.classifier[1](x)
+        return h, v
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, input_dim, da, r):
+        super().__init__()
+        self.ws1 = nn.Linear(input_dim, da, bias=False)
+        self.ws2 = nn.Linear(da, r, bias=False)
+
+    def forward(self, h):
+        attn_mat = F.softmax(self.ws2(torch.tanh(self.ws1(h))), dim=1)
+        attn_mat = attn_mat.permute(0, 2, 1)
+        return attn_mat
+
+
+class Sequence_Classifier(nn.Module):
+    def __init__(self, class_num=1, hs=512, da=512, r=8):
+        super().__init__()
+        self.attention = SelfAttention(hs, da, r)
+        self.classifier = MLP([hs * r, (hs * r + class_num) // 2, class_num])
+
+    def forward(self, x):
+        attn_mat = self.attention(x)
+        m = torch.bmm(attn_mat, x)
+        flatten = m.view(m.size()[0], -1)
+        res = self.classifier(flatten)
+        return res
+
+
+class Token_Predictor(nn.Module):
+    def __init__(self, hidden_dim=512, class_num=1):
+        super().__init__()
+        self.classifier = MLP([hidden_dim, (hidden_dim + class_num) // 2, class_num])
+
+    def forward(self, x):
+        x = self.classifier(x)
+        return x
+
+class Skip_BART(nn.Module,
+    PyTorchModelHubMixin
+):
+    def __init__(self, class_num=[180, 256], max_position_embeddings=1024, hidden_size=1024, layers=8, heads=8, ffn_dims=2048, pretrain=False):
+        super().__init__()
+        self.config = BartConfig(max_position_embeddings=max_position_embeddings,
+                                d_model=hidden_size,
+                                encoder_layers=layers,
+                                encoder_ffn_dim=ffn_dims,
+                                encoder_attention_heads=heads,
+                                decoder_layers=layers,
+                                decoder_ffn_dim=ffn_dims,
+                                decoder_attention_heads=heads
+                                )
+        self.model = ML_BART(self.config, class_num = class_num, pretrain = pretrain)
+
+    def forward(self, x_encoder, x_decoder, attn_mask_encoder=None, attn_mask_decoder=None):
+        return self.model(x_encoder, x_decoder, attn_mask_encoder, attn_mask_decoder)

trained.zip

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