def step(hparams, tokens, past=None):
lm_output = model.model(hparams=hparams,
X=tokens,
past=past,
reuse=tf.AUTO_REUSE)
logits = lm_output['logits']
presents = lm_output['present']
presents.set_shape(
model.past_shape(hparams=hparams, batch_size=None))
return {
'logits': logits,
'presents': presents,
}
def decode_all(self, hparams, tokens, past):
with tf.variable_scope(self.scope):
lm_output = model.model(hparams=hparams,
X=tokens,
past=past,
reuse=tf.AUTO_REUSE)
logits = lm_output['logits']
presents = lm_output['present']
presents.set_shape(
model.past_shape(hparams=hparams, batch_size=None))
return {
'logits': logits,
'presents': tf.concat([past, presents], axis=-2)
}
def encode(self, input, input_len, past=None):
with tf.variable_scope(self.scope, reuse=tf.AUTO_REUSE):
lm_output = model.model(hparams=self.hparam,
X=input,
past=past,
reuse=tf.AUTO_REUSE)
presents = lm_output['present']
presents.set_shape(
model.past_shape(hparams=self.hparam, batch_size=None))
target_mask = tf.sequence_mask(input_len,
maxlen=tf.shape(input)[1],
dtype=tf.float32)
target_mask = tf.expand_dims(target_mask, 2)
print(presents)
encode_out = tf.transpose(presents, perm=(0, 4, 2, 3, 1, 5))
ori_enc_shape = tf.shape(encode_out)
encode_out = tf.reshape(encode_out,
shape=(tf.shape(presents)[0],
tf.shape(presents)[4], -1))
encode_out = tf.multiply(encode_out, target_mask)
encode_out = tf.reshape(encode_out, shape=ori_enc_shape)
encode_out = tf.transpose(encode_out, perm=(0, 4, 2, 3, 1, 5))
return encode_out