def __init__(self, corpus, config):
super(HRED, self).__init__(config)
self.vocab = corpus.vocab
self.rev_vocab = corpus.rev_vocab
self.vocab_size = len(self.vocab)
self.go_id = self.rev_vocab[BOS]
self.eos_id = self.rev_vocab[EOS]
self.pad_id = self.rev_vocab[PAD]
# build model here
self.utt_encoder = ElmoUttEncoder(
config,
config.utt_cell_size,
config.dropout,
use_attn=config.utt_type == 'attn_rnn',
feat_size=1)
self.ctx_encoder = EncoderRNN(self.utt_encoder.output_size,
config.ctx_cell_size,
0.0,
config.dropout,
config.num_layer,
config.rnn_cell,
variable_lengths=False,
bidirection=config.bi_ctx_cell)
if config.bi_ctx_cell or config.num_layer > 1:
self.connector = Bi2UniConnector(config.rnn_cell, config.num_layer,
config.ctx_cell_size,
config.dec_cell_size)
else:
self.connector = IdentityConnector()
self.decoder = ElmoDecoderRNN(self.vocab_size,
config.max_dec_len,
self.utt_encoder.embed_size,
config.dec_cell_size,
self.go_id,
self.eos_id,
self.vocab,
n_layers=1,
rnn_cell=config.rnn_cell,
input_dropout_p=config.dropout,
dropout_p=config.dropout,
use_attention=config.use_attn,
embedding=self.utt_encoder.embedding,
attn_size=self.ctx_encoder.output_size,
attn_mode=config.attn_type,
use_gpu=config.use_gpu)
self.nll = criterions.NLLEntropy(self.pad_id, config)
def __init__(self, corpus, config):
super(PtrHRED, self).__init__(config, corpus)
self.ctx_encoder = EncoderRNN(self.utt_encoder.output_size,
config.ctx_cell_size,
0.0,
config.dropout,
config.num_layer,
config.rnn_cell,
variable_lengths=False,
bidirection=config.bi_ctx_cell)
if config.bi_ctx_cell or config.num_layer > 1:
self.connector = Bi2UniConnector(config.rnn_cell, config.num_layer,
config.ctx_cell_size,
config.dec_cell_size)
else:
self.connector = IdentityConnector()
self.attn_size = self.ctx_encoder.output_size
self.plain_embedding = nn.Embedding(self.vocab_size, config.embed_size)
self.decoder = DecoderPointerGen(self.vocab_size,
config.max_dec_len,
config.embed_size,
config.dec_cell_size,
self.go_id,
self.eos_id,
n_layers=1,
rnn_cell=config.rnn_cell,
input_dropout_p=config.dropout,
dropout_p=config.dropout,
attn_size=self.attn_size,
attn_mode=config.attn_type,
use_gpu=config.use_gpu,
embedding=self.plain_embedding)
self.nll_loss = criterions.NLLEntropy(self.pad_id, config)