|
import torch |
|
import torch.nn as nn |
|
import fairseq |
|
import os |
|
import hydra |
|
|
|
def load_ssl_model(cp_path): |
|
ssl_model_type = cp_path.split("/")[-1] |
|
wavlm = "WavLM" in ssl_model_type |
|
if wavlm: |
|
checkpoint = torch.load(cp_path) |
|
cfg = WavLMConfig(checkpoint['cfg']) |
|
ssl_model = WavLM(cfg) |
|
ssl_model.load_state_dict(checkpoint['model']) |
|
if 'Large' in ssl_model_type: |
|
SSL_OUT_DIM = 1024 |
|
else: |
|
SSL_OUT_DIM = 768 |
|
else: |
|
if ssl_model_type == "wav2vec_small.pt": |
|
SSL_OUT_DIM = 768 |
|
elif ssl_model_type in ["w2v_large_lv_fsh_swbd_cv.pt", "xlsr_53_56k.pt"]: |
|
SSL_OUT_DIM = 1024 |
|
else: |
|
print("*** ERROR *** SSL model type " + ssl_model_type + " not supported.") |
|
exit() |
|
model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task( |
|
[cp_path] |
|
) |
|
ssl_model = model[0] |
|
ssl_model.remove_pretraining_modules() |
|
return SSL_model(ssl_model, SSL_OUT_DIM, wavlm) |
|
|
|
class SSL_model(nn.Module): |
|
def __init__(self,ssl_model,ssl_out_dim,wavlm) -> None: |
|
super(SSL_model,self).__init__() |
|
self.ssl_model, self.ssl_out_dim = ssl_model, ssl_out_dim |
|
self.WavLM = wavlm |
|
|
|
def forward(self,batch): |
|
wav = batch['wav'] |
|
wav = wav.squeeze(1) |
|
if self.WavLM: |
|
x = self.ssl_model.extract_features(wav)[0] |
|
else: |
|
res = self.ssl_model(wav, mask=False, features_only=True) |
|
x = res["x"] |
|
return {"ssl-feature":x} |
|
def get_output_dim(self): |
|
return self.ssl_out_dim |
|
|
|
|
|
class PhonemeEncoder(nn.Module): |
|
''' |
|
PhonemeEncoder consists of an embedding layer, an LSTM layer, and a linear layer. |
|
Args: |
|
vocab_size: the size of the vocabulary |
|
hidden_dim: the size of the hidden state of the LSTM |
|
emb_dim: the size of the embedding layer |
|
out_dim: the size of the output of the linear layer |
|
n_lstm_layers: the number of LSTM layers |
|
''' |
|
def __init__(self, vocab_size, hidden_dim, emb_dim, out_dim,n_lstm_layers,with_reference=True) -> None: |
|
super().__init__() |
|
self.with_reference = with_reference |
|
self.embedding = nn.Embedding(vocab_size, emb_dim) |
|
self.encoder = nn.LSTM(emb_dim, hidden_dim, |
|
num_layers=n_lstm_layers, dropout=0.1, bidirectional=True) |
|
self.linear = nn.Sequential( |
|
nn.Linear(hidden_dim + hidden_dim*self.with_reference, out_dim), |
|
nn.ReLU() |
|
) |
|
self.out_dim = out_dim |
|
|
|
def forward(self,batch): |
|
seq = batch['phonemes'] |
|
lens = batch['phoneme_lens'] |
|
reference_seq = batch['reference'] |
|
reference_lens = batch['reference_lens'] |
|
emb = self.embedding(seq) |
|
emb = torch.nn.utils.rnn.pack_padded_sequence( |
|
emb, lens, batch_first=True, enforce_sorted=False) |
|
_, (ht, _) = self.encoder(emb) |
|
feature = ht[-1] + ht[0] |
|
if self.with_reference: |
|
if reference_seq==None or reference_lens ==None: |
|
raise ValueError("reference_batch and reference_lens should not be None when with_reference is True") |
|
reference_emb = self.embedding(reference_seq) |
|
reference_emb = torch.nn.utils.rnn.pack_padded_sequence( |
|
reference_emb, reference_lens, batch_first=True, enforce_sorted=False) |
|
_, (ht_ref, _) = self.encoder(emb) |
|
reference_feature = ht_ref[-1] + ht_ref[0] |
|
feature = self.linear(torch.cat([feature,reference_feature],1)) |
|
else: |
|
feature = self.linear(feature) |
|
return {"phoneme-feature": feature} |
|
def get_output_dim(self): |
|
return self.out_dim |
|
|
|
class DomainEmbedding(nn.Module): |
|
def __init__(self,n_domains,domain_dim) -> None: |
|
super().__init__() |
|
self.embedding = nn.Embedding(n_domains,domain_dim) |
|
self.output_dim = domain_dim |
|
def forward(self, batch): |
|
return {"domain-feature": self.embedding(batch['domains'])} |
|
def get_output_dim(self): |
|
return self.output_dim |
|
|
|
|
|
class LDConditioner(nn.Module): |
|
''' |
|
Conditions ssl output by listener embedding |
|
''' |
|
def __init__(self,input_dim, judge_dim, num_judges=None): |
|
super().__init__() |
|
self.input_dim = input_dim |
|
self.judge_dim = judge_dim |
|
self.num_judges = num_judges |
|
assert num_judges !=None |
|
self.judge_embedding = nn.Embedding(num_judges, self.judge_dim) |
|
|
|
|
|
self.decoder_rnn = nn.LSTM( |
|
input_size = self.input_dim + self.judge_dim, |
|
hidden_size = 512, |
|
num_layers = 1, |
|
batch_first = True, |
|
bidirectional = True |
|
) |
|
self.out_dim = self.decoder_rnn.hidden_size*2 |
|
|
|
def get_output_dim(self): |
|
return self.out_dim |
|
|
|
|
|
def forward(self, x, batch): |
|
judge_ids = batch['judge_id'] |
|
if 'phoneme-feature' in x.keys(): |
|
concatenated_feature = torch.cat((x['ssl-feature'], x['phoneme-feature'].unsqueeze(1).expand(-1,x['ssl-feature'].size(1) ,-1)),dim=2) |
|
else: |
|
concatenated_feature = x['ssl-feature'] |
|
if 'domain-feature' in x.keys(): |
|
concatenated_feature = torch.cat( |
|
( |
|
concatenated_feature, |
|
x['domain-feature'] |
|
.unsqueeze(1) |
|
.expand(-1, concatenated_feature.size(1), -1), |
|
), |
|
dim=2, |
|
) |
|
if judge_ids != None: |
|
concatenated_feature = torch.cat( |
|
( |
|
concatenated_feature, |
|
self.judge_embedding(judge_ids) |
|
.unsqueeze(1) |
|
.expand(-1, concatenated_feature.size(1), -1), |
|
), |
|
dim=2, |
|
) |
|
decoder_output, (h, c) = self.decoder_rnn(concatenated_feature) |
|
return decoder_output |
|
|
|
class Projection(nn.Module): |
|
def __init__(self, input_dim, hidden_dim, activation, range_clipping=False): |
|
super(Projection, self).__init__() |
|
self.range_clipping = range_clipping |
|
output_dim = 1 |
|
if range_clipping: |
|
self.proj = nn.Tanh() |
|
|
|
self.net = nn.Sequential( |
|
nn.Linear(input_dim, hidden_dim), |
|
activation, |
|
nn.Dropout(0.3), |
|
nn.Linear(hidden_dim, output_dim), |
|
) |
|
self.output_dim = output_dim |
|
|
|
def forward(self, x, batch): |
|
output = self.net(x) |
|
|
|
|
|
if self.range_clipping: |
|
return self.proj(output) * 2.0 + 3 |
|
else: |
|
return output |
|
def get_output_dim(self): |
|
return self.output_dim |
|
|
