def test_copy_gen_outp_has_no_prob_of_pad(self): for params, init_case in itertools.product(self.PARAMS, self.INIT_CASES): cgen = CopyGenerator(**init_case) dummy_in = self.dummy_inputs(params, init_case) res = cgen(*dummy_in) self.assertTrue(res[:, init_case["pad_idx"]].allclose( torch.tensor(0.0)))
def test_copy_gen_forward_shape(self): for params, init_case in itertools.product(self.PARAMS, self.INIT_CASES): cgen = CopyGenerator(**init_case) dummy_in = self.dummy_inputs(params, init_case) res = cgen(*dummy_in) expected_shape = self.expected_shape(params, init_case) self.assertEqual(res.shape, expected_shape, init_case.__str__())
def test_copy_gen_trainable_params_update(self): for params, init_case in itertools.product( self.PARAMS, self.INIT_CASES): cgen = CopyGenerator(**init_case) trainable_params = {n: p for n, p in cgen.named_parameters() if p.requires_grad} assert len(trainable_params) > 0 # sanity check old_weights = deepcopy(trainable_params) dummy_in = self.dummy_inputs(params, init_case) res = cgen(*dummy_in) pretend_loss = res.sum() pretend_loss.backward() dummy_optim = torch.optim.SGD(trainable_params.values(), 1) dummy_optim.step() for param_name in old_weights.keys(): self.assertTrue( trainable_params[param_name] .ne(old_weights[param_name]).any(), param_name + " " + init_case.__str__())
def test_copy_gen_trainable_params_update(self): for params, init_case in itertools.product( self.PARAMS, self.INIT_CASES): cgen = CopyGenerator(**init_case) trainable_params = {n: p for n, p in cgen.named_parameters() if p.requires_grad} assert len(trainable_params) > 0 # sanity check old_weights = deepcopy(trainable_params) dummy_in = self.dummy_inputs(params, init_case) res = cgen(*dummy_in) pretend_loss = res.sum() pretend_loss.backward() dummy_optim = torch.optim.SGD(trainable_params.values(), 1) dummy_optim.step() for param_name in old_weights.keys(): self.assertTrue( trainable_params[param_name] .ne(old_weights[param_name]).any(), param_name + " " + init_case.__str__())
def build_tm_model(opt, dicts): onmt.constants.neg_log_sigma1 = opt.neg_log_sigma1 onmt.constants.neg_log_sigma2 = opt.neg_log_sigma2 onmt.constants.prior_pi = opt.prior_pi # BUILD POSITIONAL ENCODING if opt.time == 'positional_encoding': positional_encoder = PositionalEncoding(opt.model_size, len_max=MAX_LEN) else: raise NotImplementedError # BUILD GENERATOR if opt.copy_generator: generators = [ CopyGenerator(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] else: generators = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] # BUILD EMBEDDINGS if 'src' in dicts: embedding_src = nn.Embedding(dicts['src'].size(), opt.model_size, padding_idx=onmt.constants.PAD) else: embedding_src = None if opt.join_embedding and embedding_src is not None: embedding_tgt = embedding_src print("* Joining the weights of encoder and decoder word embeddings") else: embedding_tgt = nn.Embedding(dicts['tgt'].size(), opt.model_size, padding_idx=onmt.constants.PAD) if opt.use_language_embedding: print("* Create language embeddings with %d languages" % len(dicts['langs'])) language_embeddings = nn.Embedding(len(dicts['langs']), opt.model_size) else: language_embeddings = None if opt.encoder_type == "text": encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = RelativeTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) if opt.reconstruct: rev_decoder = RelativeTransformerDecoder( opt, embedding_src, None, language_embeddings=language_embeddings) rev_generator = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['src'].size(), fix_norm=opt.fix_norm_output_embedding) ] rev_generator = nn.ModuleList(rev_generator) else: rev_decoder = None rev_generator = None model = BayesianTransformer(encoder, decoder, generator, rev_decoder, rev_generator, mirror=opt.mirror_loss) if opt.tie_weights: print("* Joining the weights of decoder input and output embeddings") model.tie_weights() return model
def build_tm_model(opt, dicts): # BUILD POSITIONAL ENCODING if opt.time == 'positional_encoding': positional_encoder = PositionalEncoding(opt.model_size, len_max=MAX_LEN) else: raise NotImplementedError if opt.reconstruct: # reconstruction is only compatible assert opt.model == 'relative_transformer' assert opt.encoder_type == 'text' # BUILD GENERATOR if opt.copy_generator: if opt.nce_noise > 0: print("[INFO] Copy generator overrides NCE.") opt.nce = False opt.nce_noise = 0 generators = [ CopyGenerator(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] elif opt.nce_noise > 0: from onmt.modules.nce.nce_linear import NCELinear from onmt.modules.nce.nce_utils import build_unigram_noise noise_distribution = build_unigram_noise( torch.FloatTensor(list(dicts['tgt'].frequencies.values()))) generator = NCELinear(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding, noise_distribution=noise_distribution, noise_ratio=opt.nce_noise) generators = [generator] else: generators = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] # BUILD EMBEDDINGS if 'src' in dicts: embedding_src = nn.Embedding(dicts['src'].size(), opt.model_size, padding_idx=onmt.constants.PAD) else: embedding_src = None if opt.join_embedding and embedding_src is not None: embedding_tgt = embedding_src print("* Joining the weights of encoder and decoder word embeddings") else: embedding_tgt = nn.Embedding(dicts['tgt'].size(), opt.model_size, padding_idx=onmt.constants.PAD) if opt.use_language_embedding: print("* Create language embeddings with %d languages" % len(dicts['langs'])) language_embeddings = nn.Embedding(len(dicts['langs']), opt.model_size) else: language_embeddings = None if opt.ctc_loss != 0: generators.append( onmt.modules.base_seq2seq.Generator(opt.model_size, dicts['tgt'].size() + 1)) if opt.model in ['conformer', 'speech_transformer', 'hybrid_transformer']: onmt.constants.init_value = opt.param_init from onmt.models.speech_recognizer.relative_transformer import \ SpeechTransformerEncoder, SpeechTransformerDecoder if opt.model == 'conformer': from onmt.models.speech_recognizer.conformer import ConformerEncoder, Conformer from onmt.models.speech_recognizer.lstm import SpeechLSTMDecoder opt.cnn_downsampling = True # force this bool to have masking at decoder to be corrected encoder = ConformerEncoder(opt, None, None, 'audio') decoder = SpeechLSTMDecoder( opt, embedding_tgt, language_embeddings=language_embeddings) model = Conformer(encoder, decoder, nn.ModuleList(generators), ctc=opt.ctc_loss > 0.0) elif opt.model == 'hybrid_transformer': from onmt.models.speech_recognizer.lstm import SpeechLSTMDecoder, SpeechLSTMEncoder, SpeechLSTMSeq2Seq encoder = SpeechTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = SpeechLSTMDecoder( opt, embedding_tgt, language_embeddings=language_embeddings) model = SpeechLSTMSeq2Seq(encoder, decoder, nn.ModuleList(generators), ctc=opt.ctc_loss > 0.0) else: encoder = SpeechTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = SpeechTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss, ctc=opt.ctc_loss > 0.0) # If we use the multilingual model and weights are partitioned: if opt.multilingual_partitioned_weights: # this is basically the language embeddings factor_embeddings = nn.Embedding(len(dicts['langs']), opt.mpw_factor_size) encoder.factor_embeddings = factor_embeddings decoder.factor_embeddings = factor_embeddings elif opt.model in ["LSTM", 'lstm']: # print("LSTM") onmt.constants.init_value = opt.param_init from onmt.models.speech_recognizer.lstm import SpeechLSTMDecoder, SpeechLSTMEncoder, SpeechLSTMSeq2Seq encoder = SpeechLSTMEncoder(opt, None, opt.encoder_type) decoder = SpeechLSTMDecoder(opt, embedding_tgt, language_embeddings=language_embeddings) model = SpeechLSTMSeq2Seq(encoder, decoder, nn.ModuleList(generators), ctc=opt.ctc_loss > 0.0) elif opt.model in ['multilingual_translator', 'translator']: onmt.constants.init_value = opt.param_init from onmt.models.multilingual_translator.relative_transformer import \ RelativeTransformerEncoder, RelativeTransformerDecoder encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss) elif opt.model in ['transformer', 'stochastic_transformer']: onmt.constants.init_value = opt.param_init if opt.encoder_type == "text": encoder = TransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = TransformerEncoder(opt, None, positional_encoder, opt.encoder_type) elif opt.encoder_type == "mix": text_encoder = TransformerEncoder( opt, embedding_src, positional_encoder, "text", language_embeddings=language_embeddings) audio_encoder = TransformerEncoder(opt, None, positional_encoder, "audio") encoder = MixedEncoder(text_encoder, audio_encoder) else: print("Unknown encoder type:", opt.encoder_type) exit(-1) decoder = TransformerDecoder(opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, nn.ModuleList(generators), mirror=opt.mirror_loss) elif opt.model == 'relative_transformer': from onmt.models.relative_transformer import \ RelativeTransformerEncoder, RelativeTransformerDecoder if opt.encoder_type == "text": encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = RelativeTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) if opt.reconstruct: rev_decoder = RelativeTransformerDecoder( opt, embedding_src, None, language_embeddings=language_embeddings) rev_generator = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['src'].size(), fix_norm=opt.fix_norm_output_embedding) ] rev_generator = nn.ModuleList(rev_generator) else: rev_decoder = None rev_generator = None model = RelativeTransformer(encoder, decoder, generator, rev_decoder, rev_generator, mirror=opt.mirror_loss) elif opt.model == 'distance_transformer': # from onmt.models.relative_transformer import RelativeTransformerDecoder, RelativeTransformer from onmt.models.distance_transformer import DistanceTransformerEncoder if opt.encoder_type == "text": encoder = DistanceTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = DistanceTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = DistanceTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'universal_transformer': from onmt.models.universal_transformer import UniversalTransformerDecoder, UniversalTransformerEncoder generator = nn.ModuleList(generators) if opt.encoder_type == "text": encoder = UniversalTransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = UniversalTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = UniversalTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'relative_universal_transformer': from onmt.models.relative_universal_transformer import \ RelativeUniversalTransformerEncoder, RelativeUniversalTransformerDecoder generator = nn.ModuleList(generators) if opt.encoder_type == "text": encoder = RelativeUniversalTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = RelativeUniversalTransformerDecoder( opt, None, None, opt.encoder_type) decoder = RelativeUniversalTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'relative_unified_transformer': from onmt.models.relative_unified_transformer import RelativeUnifiedTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = RelativeUnifiedTransformer( opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings) elif opt.model == 'memory_transformer': from onmt.models.memory_transformer import MemoryTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = MemoryTransformer(opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings, dictionary=dicts['tgt']) else: raise NotImplementedError if opt.tie_weights: print("* Joining the weights of decoder input and output embeddings") model.tie_weights() return model
def build_tm_model(opt, dicts): # BUILD POSITIONAL ENCODING if opt.time == 'positional_encoding': positional_encoder = PositionalEncoding(opt.model_size, len_max=MAX_LEN) else: raise NotImplementedError # BUILD GENERATOR if opt.copy_generator: generators = [ CopyGenerator(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] else: generators = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] # BUILD EMBEDDINGS if 'src' in dicts: embedding_src = nn.Embedding(dicts['src'].size(), opt.model_size, padding_idx=onmt.constants.PAD) else: embedding_src = None if opt.join_embedding and embedding_src is not None: embedding_tgt = embedding_src print("* Joining the weights of encoder and decoder word embeddings") else: embedding_tgt = nn.Embedding(dicts['tgt'].size(), opt.model_size, padding_idx=onmt.constants.PAD) if opt.use_language_embedding: print("* Create language embeddings with %d languages" % len(dicts['langs'])) language_embeddings = nn.Embedding(len(dicts['langs']), opt.model_size) else: language_embeddings = None if opt.ctc_loss != 0: generators.append( onmt.modules.base_seq2seq.Generator(opt.model_size, dicts['tgt'].size() + 1)) if opt.model in ['transformer', 'stochastic_transformer']: onmt.constants.init_value = opt.param_init if opt.encoder_type == "text": encoder = TransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = TransformerEncoder(opt, None, positional_encoder, opt.encoder_type) elif opt.encoder_type == "mix": text_encoder = TransformerEncoder( opt, embedding_src, positional_encoder, "text", language_embeddings=language_embeddings) audio_encoder = TransformerEncoder(opt, None, positional_encoder, "audio") encoder = MixedEncoder(text_encoder, audio_encoder) else: print("Unknown encoder type:", opt.encoder_type) exit(-1) decoder = TransformerDecoder(opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, nn.ModuleList(generators), mirror=opt.mirror_loss) elif opt.model == 'relative_transformer': from onmt.models.relative_transformer import RelativeTransformerEncoder, RelativeTransformerDecoder, \ RelativeTransformer if opt.encoder_type == "text": encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = RelativeTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'distance_transformer': from onmt.models.relative_transformer import RelativeTransformerDecoder, RelativeTransformer from onmt.models.distance_transformer import DistanceTransformerEncoder if opt.encoder_type == "text": encoder = DistanceTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = DistanceTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'unified_transformer': from onmt.models.unified_transformer import UnifiedTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = UnifiedTransformer(opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings) elif opt.model == 'relative_unified_transformer': from onmt.models.relative_unified_transformer import RelativeUnifiedTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = RelativeUnifiedTransformer( opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings) elif opt.model == 'memory_transformer': from onmt.models.memory_transformer import MemoryTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = MemoryTransformer(opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings, dictionary=dicts['tgt']) else: raise NotImplementedError if opt.tie_weights: print("* Joining the weights of decoder input and output embeddings") model.tie_weights() return model
def build_tm_model(opt, dicts): # BUILD POSITIONAL ENCODING if opt.time == 'positional_encoding': positional_encoder = PositionalEncoding(opt.model_size, len_max=MAX_LEN) else: raise NotImplementedError # temporary fix for speech autoencoder if opt.reconstruct: # reconstruction is only compatible assert opt.model == 'relative_transformer' assert opt.encoder_type == 'text' # BUILD GENERATOR if opt.copy_generator: if opt.nce_noise > 0: print("[INFO] Copy generator overrides NCE.") opt.nce = False opt.nce_noise = 0 generators = [ CopyGenerator(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] elif opt.nce_noise > 0: from onmt.modules.nce.nce_linear import NCELinear from onmt.modules.nce.nce_utils import build_unigram_noise noise_distribution = build_unigram_noise( torch.FloatTensor(list(dicts['tgt'].frequencies.values()))) generator = NCELinear(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding, noise_distribution=noise_distribution, noise_ratio=opt.nce_noise) generators = [generator] else: if "tgt" in dicts: generators = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] # BUILD EMBEDDINGS if 'src' in dicts: embedding_src = nn.Embedding(dicts['src'].size(), opt.model_size, padding_idx=onmt.constants.PAD) else: embedding_src = None if opt.join_embedding and embedding_src is not None: embedding_tgt = embedding_src print("* Joining the weights of encoder and decoder word embeddings") elif 'tgt' in dicts: embedding_tgt = nn.Embedding(dicts['tgt'].size(), opt.model_size, padding_idx=onmt.constants.PAD) # done else: embedding_tgt = None if opt.use_language_embedding: # done print("* Create language embeddings with %d languages" % len(dicts['langs'])) language_embeddings = nn.Embedding(len(dicts['langs']), opt.model_size) else: language_embeddings = None if opt.ctc_loss != 0: generators.append( onmt.modules.base_seq2seq.Generator(opt.model_size, dicts['tgt'].size() + 1)) if opt.model in ['conformer', 'speech_transformer']: onmt.constants.init_value = opt.param_init from onmt.models.speech_recognizer.relative_transformer import \ SpeechTransformerEncoder, SpeechTransformerDecoder if opt.model == 'conformer': from onmt.models.speech_recognizer.conformer import ConformerEncoder opt.cnn_downsampling = True # force this bool to have masking at decoder to be corrected encoder = ConformerEncoder(opt, None, None, 'audio') else: encoder = SpeechTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = SpeechTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss) elif opt.model == "LSTM_v2": onmt.constants.init_value = opt.param_init from onmt.models.speech_recognizer.lstm_v2 import SpeechLSTMDecoder, SpeechLSTMEncoder, SpeechLSTMSeq2Seq encoder = SpeechLSTMEncoder(opt, None, opt.encoder_type) decoder = SpeechLSTMDecoder(opt, embedding_tgt, language_embeddings=language_embeddings) model = SpeechLSTMSeq2Seq(encoder, decoder, nn.ModuleList(generators)) # print(model) # sys.exit() elif opt.model == "LSTM": onmt.constants.init_value = opt.param_init from onmt.models.LSTM_based import SpeechLSTMDecoder, SpeechLSTMEncoder, SpeechLSTMSeq2Seq encoder = SpeechLSTMEncoder(opt, None, opt.encoder_type) decoder = SpeechLSTMDecoder(opt, embedding_tgt, language_embeddings=language_embeddings) model = SpeechLSTMSeq2Seq(encoder, decoder, nn.ModuleList(generators)) elif opt.model in ['multilingual_translator', 'translator']: onmt.constants.init_value = opt.param_init from onmt.models.multilingual_translator.relative_transformer import \ RelativeTransformerEncoder, RelativeTransformerDecoder encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss) elif opt.model == "speech_ae" or opt.model == "speech2speech": onmt.constants.init_value = opt.param_init from onmt.models.speech_metamorphosis.speechae import SpeechLSTMEncoder from onmt.models.speech_metamorphosis.speechae import TacotronDecoder, SpeechAE encoder = SpeechLSTMEncoder(opt, None, opt.encoder_type) decoder = TacotronDecoder(opt) model = SpeechAE(encoder, decoder, opt) print("Create speech autoencoder successfully") elif opt.model == "speech_FN": onmt.constants.init_value = opt.param_init from onmt.models.speech_metamorphosis.speechae import SpeechLSTMEncoder, LatentDiscrinator from onmt.models.speech_metamorphosis.speechae import TacotronDecoder, SpeechAE encoder = SpeechLSTMEncoder(opt, None, opt.encoder_type) decoder = TacotronDecoder(opt, accent_emdedding=language_embeddings) model_ae = SpeechAE(encoder, decoder, opt) lat_dis = LatentDiscrinator(opt, hidden_size=128, output_size=len(dicts['langs'])) model = (model_ae, lat_dis) print("Create speech fader network successfully") elif opt.model in ['transformer', 'stochastic_transformer']: onmt.constants.init_value = opt.param_init if opt.encoder_type == "text": encoder = TransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = TransformerEncoder(opt, None, positional_encoder, opt.encoder_type) elif opt.encoder_type == "mix": text_encoder = TransformerEncoder( opt, embedding_src, positional_encoder, "text", language_embeddings=language_embeddings) audio_encoder = TransformerEncoder(opt, None, positional_encoder, "audio") encoder = MixedEncoder(text_encoder, audio_encoder) else: print("Unknown encoder type:", opt.encoder_type) exit(-1) decoder = TransformerDecoder(opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, nn.ModuleList(generators), mirror=opt.mirror_loss) elif opt.model == 'relative_transformer': from onmt.models.relative_transformer import \ RelativeTransformerEncoder, RelativeTransformerDecoder if opt.encoder_type == "text": encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = RelativeTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) if opt.reconstruct: rev_decoder = RelativeTransformerDecoder( opt, embedding_src, None, language_embeddings=language_embeddings) rev_generator = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['src'].size(), fix_norm=opt.fix_norm_output_embedding) ] rev_generator = nn.ModuleList(rev_generator) else: rev_decoder = None rev_generator = None model = RelativeTransformer(encoder, decoder, generator, rev_decoder, rev_generator, mirror=opt.mirror_loss) elif opt.model == 'universal_transformer': from onmt.models.universal_transformer import UniversalTransformerDecoder, UniversalTransformerEncoder generator = nn.ModuleList(generators) if opt.encoder_type == "text": encoder = UniversalTransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = UniversalTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = UniversalTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'relative_universal_transformer': from onmt.models.relative_universal_transformer import \ RelativeUniversalTransformerEncoder, RelativeUniversalTransformerDecoder generator = nn.ModuleList(generators) if opt.encoder_type == "text": encoder = RelativeUniversalTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = RelativeUniversalTransformerDecoder( opt, None, None, opt.encoder_type) decoder = RelativeUniversalTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'relative_unified_transformer': from onmt.models.relative_unified_transformer import RelativeUnifiedTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = RelativeUnifiedTransformer( opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings) elif opt.model == 'memory_transformer': from onmt.models.memory_transformer import MemoryTransformer if opt.encoder_type == "audio": raise NotImplementedError generator = nn.ModuleList(generators) model = MemoryTransformer(opt, embedding_src, embedding_tgt, generator, positional_encoder, language_embeddings=language_embeddings, dictionary=dicts['tgt']) else: raise NotImplementedError if opt.tie_weights: print("* Joining the weights of decoder input and output embeddings") model.tie_weights() if opt.tie_weights_lid: print("* Joining the weights of output lid and language embeddings") model.tie_weights_lid() return model
def build_tm_model(opt, dicts): onmt.constants = add_tokenidx(opt, onmt.constants, dicts) # BUILD POSITIONAL ENCODING if opt.time == 'positional_encoding': positional_encoder = PositionalEncoding(opt.model_size, len_max=MAX_LEN) else: raise NotImplementedError if opt.reconstruct: # reconstruction is only compatible assert opt.model == 'relative_transformer' assert opt.encoder_type == 'text' # BUILD GENERATOR if opt.copy_generator: if opt.nce_noise > 0: print("[INFO] Copy generator overrides NCE.") opt.nce = False opt.nce_noise = 0 generators = [ CopyGenerator(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] elif opt.nce_noise > 0: from onmt.modules.nce.nce_linear import NCELinear from onmt.modules.nce.nce_utils import build_unigram_noise noise_distribution = build_unigram_noise( torch.FloatTensor(list(dicts['tgt'].frequencies.values()))) generator = NCELinear(opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding, noise_distribution=noise_distribution, noise_ratio=opt.nce_noise) generators = [generator] else: generators = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['tgt'].size(), fix_norm=opt.fix_norm_output_embedding) ] # BUILD EMBEDDINGS if 'src' in dicts: if (not hasattr(opt, "enc_pretrained_model")) or ( not opt.enc_pretrained_model): embedding_src = nn.Embedding(dicts['src'].size(), opt.model_size, padding_idx=onmt.constants.SRC_PAD) else: embedding_src = None if opt.join_embedding and embedding_src is not None: embedding_tgt = embedding_src print("* Joining the weights of encoder and decoder word embeddings") elif not opt.dec_pretrained_model: embedding_tgt = nn.Embedding(dicts['tgt'].size(), opt.model_size, padding_idx=onmt.constants.TGT_PAD) else: assert opt.model == "pretrain_transformer" embedding_tgt = None if opt.use_language_embedding: print("* Create language embeddings with %d languages" % len(dicts['langs'])) language_embeddings = nn.Embedding(len(dicts['langs']), opt.model_size) else: language_embeddings = None if opt.ctc_loss != 0: generators.append( onmt.modules.base_seq2seq.Generator(opt.model_size, dicts['tgt'].size() + 1)) if opt.model in ['conformer', 'speech_transformer', 'hybrid_transformer']: onmt.constants.init_value = opt.param_init from onmt.models.speech_recognizer.relative_transformer import \ SpeechTransformerEncoder, SpeechTransformerDecoder if opt.model == 'conformer': from onmt.models.speech_recognizer.conformer import ConformerEncoder, Conformer from onmt.models.speech_recognizer.lstm import SpeechLSTMDecoder opt.cnn_downsampling = True # force this bool to have masking at decoder to be corrected encoder = ConformerEncoder(opt, None, None, 'audio') # decoder = SpeechLSTMDecoder(opt, embedding_tgt, language_embeddings=language_embeddings) decoder = SpeechTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) # model = Conformer(encoder, decoder, nn.ModuleList(generators), ctc=opt.ctc_loss > 0.0) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss, ctc=opt.ctc_loss > 0.0) elif opt.model == 'hybrid_transformer': from onmt.models.speech_recognizer.lstm import SpeechLSTMDecoder, SpeechLSTMEncoder, SpeechLSTMSeq2Seq encoder = SpeechTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = SpeechLSTMDecoder( opt, embedding_tgt, language_embeddings=language_embeddings) model = SpeechLSTMSeq2Seq(encoder, decoder, nn.ModuleList(generators), ctc=opt.ctc_loss > 0.0) else: encoder = SpeechTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = SpeechTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss, ctc=opt.ctc_loss > 0.0) # If we use the multilingual model and weights are partitioned: if opt.multilingual_partitioned_weights: # this is basically the language embeddings factor_embeddings = nn.Embedding(len(dicts['langs']), opt.mpw_factor_size) encoder.factor_embeddings = factor_embeddings decoder.factor_embeddings = factor_embeddings elif opt.model in ["LSTM", 'lstm']: # print("LSTM") onmt.constants.init_value = opt.param_init from onmt.models.speech_recognizer.lstm import SpeechLSTMDecoder, SpeechLSTMEncoder, SpeechLSTMSeq2Seq encoder = SpeechLSTMEncoder(opt, None, opt.encoder_type) decoder = SpeechLSTMDecoder(opt, embedding_tgt, language_embeddings=language_embeddings) model = SpeechLSTMSeq2Seq(encoder, decoder, nn.ModuleList(generators), ctc=opt.ctc_loss > 0.0) elif opt.model in ['multilingual_translator', 'translator']: onmt.constants.init_value = opt.param_init from onmt.models.multilingual_translator.relative_transformer import \ RelativeTransformerEncoder, RelativeTransformerDecoder encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) model = RelativeTransformer(encoder, decoder, nn.ModuleList(generators), None, None, mirror=opt.mirror_loss) elif opt.model in ['transformer', 'stochastic_transformer']: onmt.constants.init_value = opt.param_init if opt.encoder_type == "text": encoder = TransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = TransformerEncoder(opt, None, positional_encoder, opt.encoder_type) elif opt.encoder_type == "mix": text_encoder = TransformerEncoder( opt, embedding_src, positional_encoder, "text", language_embeddings=language_embeddings) audio_encoder = TransformerEncoder(opt, None, positional_encoder, "audio") encoder = MixedEncoder(text_encoder, audio_encoder) else: print("Unknown encoder type:", opt.encoder_type) exit(-1) decoder = TransformerDecoder(opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, nn.ModuleList(generators), mirror=opt.mirror_loss) elif opt.model == 'relative_transformer': from onmt.models.relative_transformer import \ RelativeTransformerEncoder, RelativeTransformerDecoder if opt.encoder_type == "text": encoder = RelativeTransformerEncoder( opt, embedding_src, None, opt.encoder_type, language_embeddings=language_embeddings) if opt.encoder_type == "audio": # raise NotImplementedError encoder = RelativeTransformerEncoder( opt, None, None, encoder_type=opt.encoder_type, language_embeddings=language_embeddings) generator = nn.ModuleList(generators) decoder = RelativeTransformerDecoder( opt, embedding_tgt, None, language_embeddings=language_embeddings) if opt.reconstruct: rev_decoder = RelativeTransformerDecoder( opt, embedding_src, None, language_embeddings=language_embeddings) rev_generator = [ onmt.modules.base_seq2seq.Generator( opt.model_size, dicts['src'].size(), fix_norm=opt.fix_norm_output_embedding) ] rev_generator = nn.ModuleList(rev_generator) else: rev_decoder = None rev_generator = None model = RelativeTransformer(encoder, decoder, generator, rev_decoder, rev_generator, mirror=opt.mirror_loss) elif opt.model == 'universal_transformer': from onmt.legacy.old_models.universal_transformer import UniversalTransformerDecoder, UniversalTransformerEncoder generator = nn.ModuleList(generators) if opt.encoder_type == "text": encoder = UniversalTransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) elif opt.encoder_type == "audio": encoder = UniversalTransformerEncoder(opt, None, positional_encoder, opt.encoder_type) decoder = UniversalTransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) model = Transformer(encoder, decoder, generator, mirror=opt.mirror_loss) elif opt.model == 'pretrain_transformer': assert (opt.enc_pretrained_model or opt.dec_pretrained_model) from onmt.models.pretrain_transformer import PretrainTransformer print(f"pos_emb_type: {opt.pos_emb_type}") print(f"max_pos_length: {opt.max_pos_length }") print( f"Share position embeddings cross heads: {not opt.diff_head_pos}") print() if opt.enc_pretrained_model: print("* Build encoder with enc_pretrained_model: {}".format( opt.enc_pretrained_model)) if opt.enc_pretrained_model == "bert": from pretrain_module.configuration_bert import BertConfig from pretrain_module.modeling_bert import BertModel enc_bert_config = BertConfig.from_json_file(opt.enc_config_file) encoder = BertModel( enc_bert_config, bert_word_dropout=opt.enc_pretrain_word_dropout, bert_emb_dropout=opt.enc_pretrain_emb_dropout, bert_atten_dropout=opt.enc_pretrain_attn_dropout, bert_hidden_dropout=opt.enc_pretrain_hidden_dropout, bert_hidden_size=opt.enc_pretrain_hidden_size, is_decoder=False, before_plm_output_ln=opt.before_enc_output_ln, gradient_checkpointing=opt.enc_gradient_checkpointing, max_pos_len=opt.max_pos_length, diff_head_pos=opt.diff_head_pos, pos_emb_type=opt.pos_emb_type, ) elif opt.enc_pretrained_model == "roberta": from pretrain_module.configuration_roberta import RobertaConfig from pretrain_module.modeling_roberta import RobertaModel enc_roberta_config = RobertaConfig.from_json_file( opt.enc_config_file) encoder = RobertaModel( enc_roberta_config, bert_word_dropout=opt.enc_pretrain_word_dropout, bert_emb_dropout=opt.enc_pretrain_emb_dropout, bert_atten_dropout=opt.enc_pretrain_attn_dropout, bert_hidden_dropout=opt.enc_pretrain_hidden_dropout, bert_hidden_size=opt.enc_pretrain_hidden_size, is_decoder=False, before_plm_output_ln=opt.before_enc_output_ln, gradient_checkpointing=opt.enc_gradient_checkpointing, max_pos_len=opt.max_pos_length, diff_head_pos=opt.diff_head_pos, pos_emb_type=opt.pos_emb_type, ) elif not opt.enc_pretrained_model: print(" Encoder is not from pretrained model") encoder = TransformerEncoder( opt, embedding_src, positional_encoder, opt.encoder_type, language_embeddings=language_embeddings) else: print("Warning: only bert and roberta are implemented for encoder") exit(-1) if opt.load_from or not opt.enc_state_dict: if opt.verbose: print(" No weights loading from {} for encoder".format( opt.enc_pretrained_model)) elif opt.enc_pretrained_model: print(" Loading weights for encoder from: \n", opt.enc_state_dict) enc_model_state_dict = torch.load(opt.enc_state_dict, map_location="cpu") encoder.from_pretrained(state_dict=enc_model_state_dict, model=encoder, output_loading_info=opt.verbose, model_prefix=opt.enc_pretrained_model) if opt.dec_pretrained_model: print("* Build decoder with dec_pretrained_model: {}".format( opt.dec_pretrained_model)) if opt.dec_pretrained_model == "bert": if opt.enc_pretrained_model != "bert": from pretrain_module.configuration_bert import BertConfig from pretrain_module.modeling_bert import BertModel dec_bert_config = BertConfig.from_json_file(opt.dec_config_file) decoder = BertModel( dec_bert_config, bert_word_dropout=opt.dec_pretrain_word_dropout, bert_emb_dropout=opt.dec_pretrain_emb_dropout, bert_atten_dropout=opt.dec_pretrain_attn_dropout, bert_hidden_dropout=opt.dec_pretrain_hidden_dropout, bert_hidden_size=opt.dec_pretrain_hidden_size, is_decoder=True, gradient_checkpointing=opt.dec_gradient_checkpointing, max_pos_len=opt.max_pos_length, diff_head_pos=opt.diff_head_pos, pos_emb_type=opt.pos_emb_type, ) elif opt.dec_pretrained_model == "roberta": if opt.enc_pretrained_model != "roberta": from pretrain_module.configuration_roberta import RobertaConfig from pretrain_module.modeling_roberta import RobertaModel dec_roberta_config = RobertaConfig.from_json_file( opt.dec_config_file) decoder = RobertaModel( dec_roberta_config, bert_word_dropout=opt.dec_pretrain_word_dropout, bert_emb_dropout=opt.dec_pretrain_emb_dropout, bert_atten_dropout=opt.dec_pretrain_attn_dropout, bert_hidden_dropout=opt.dec_pretrain_hidden_dropout, bert_hidden_size=opt.dec_pretrain_hidden_size, is_decoder=True, gradient_checkpointing=opt.dec_gradient_checkpointing, max_pos_len=opt.max_pos_length, diff_head_pos=opt.diff_head_pos, pos_emb_type=opt.pos_emb_type, ) elif not opt.dec_pretrained_model: print(" Decoder is not from pretrained model") decoder = TransformerDecoder( opt, embedding_tgt, positional_encoder, language_embeddings=language_embeddings) else: print("Warning: only bert and roberta are implemented for decoder") exit(-1) if opt.load_from or not opt.dec_state_dict: if opt.verbose: print(" No weights loading from {} for decoder".format( opt.dec_pretrained_model)) elif opt.enc_pretrained_model: print(" Loading weights for decoder from: \n", opt.dec_state_dict) dec_model_state_dict = torch.load(opt.dec_state_dict, map_location="cpu") decoder.from_pretrained(state_dict=dec_model_state_dict, model=decoder, output_loading_info=opt.verbose, model_prefix=opt.dec_pretrained_model) encoder.enc_pretrained_model = opt.enc_pretrained_model decoder.dec_pretrained_model = opt.dec_pretrained_model encoder.input_type = opt.encoder_type model = PretrainTransformer(encoder, decoder, nn.ModuleList(generators)) else: raise NotImplementedError if opt.tie_weights: print("* Joining the weights of decoder input and output embeddings") model.tie_weights() return model