class Pronunciation2Spelling(nn.Module): def __init__(self, enc_config, dec_config, first_train): super(Pronunciation2Spelling, self).__init__() # KoELECTRA-base-v3 self.electra = ElectraModel(enc_config) if first_train: self.electra.load_state_dict( torch.load('../pretrained/electra_pretrained_base')) self.embedding = self.electra.get_input_embeddings() if enc_config.embedding_size != dec_config.hidden_size: self.embedding_projection = nn.Linear(enc_config.embedding_size, dec_config.hidden_size) self.decoders = Decoders(dec_config) self.dense = nn.Linear(dec_config.hidden_size, dec_config.trg_vocab_size) self.padding_idx = dec_config.padding_idx def forward(self, enc_ids, dec_ids): dec_embeddings = self.embedding(dec_ids) if hasattr(self, 'embedding_projection'): dec_embeddings = self.embedding_projection(dec_embeddings) enc_outputs = self.electra(enc_ids).last_hidden_state dec_outputs, _, _ = self.decoders(enc_ids, enc_outputs, dec_ids, dec_embeddings) model_output = self.dense(dec_outputs) return model_output
class Pronunciation2Spelling(nn.Module): def __init__(self, enc_config, dec_config): super(Pronunciation2Spelling, self).__init__() self.encoders = ElectraModel(enc_config) self.embedding = self.encoders.get_input_embeddings() if enc_config.embedding_size != dec_config.hidden_size: self.embedding_projection = nn.Linear(enc_config.embedding_size, dec_config.hidden_size) self.decoders = Decoders(dec_config) self.dense = nn.Linear(dec_config.hidden_size, dec_config.trg_vocab_size) self.padding_idx = dec_config.padding_idx def forward(self, enc_ids, dec_ids): dec_embeddings = self.embedding(dec_ids) if hasattr(self, 'embedding_projection'): dec_embeddings = self.embedding_projection(dec_embeddings) enc_outputs = self.encoders(enc_ids).last_hidden_state dec_outputs, _, _ = self.decoders(enc_ids, enc_outputs, dec_ids, dec_embeddings) model_output = self.dense(dec_outputs) return model_output