Skip-BART / model.py
RS2002's picture
Upload 2 files
3afa065 verified
raw
history blame
6.92 kB
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)