|
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 = self.encoder(x_encoder)
|
|
|
|
if self.pretrain:
|
|
|
|
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 = 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)
|
|
|