| 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_emb2 = 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 = self.encoder(x_encoder)
|
|
|
| if self.pretrain:
|
|
|
| emb_decoder = self.encoder(x_decoder)
|
| else:
|
| emb_decoder = torch.concatenate(
|
| [self.decoder_emb2[0](x_decoder[..., 0]), self.decoder_emb2[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 = 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)
|
|
|
