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