def run_decoder(self,
sequences: mx.nd.NDArray,
bucket_key: Tuple[int, int],
model_state: 'ModelState') -> Tuple[mx.nd.NDArray, mx.nd.NDArray, 'ModelState']:
"""
Runs forward pass of the single-step decoder.
:return: Probability distribution over next word, attention scores, updated model state.
"""
batch = mx.io.DataBatch(
data=[sequences.as_in_context(self.context)] + model_state.states,
label=None,
bucket_key=bucket_key,
provide_data=self._get_decoder_data_shapes(bucket_key))
self.decoder_module.forward(data_batch=batch, is_train=False)
probs, attention_probs, *model_state.states = self.decoder_module.get_outputs()
return probs, attention_probs, model_state
def run_decoder(self,
encoded_source: mx.nd.NDArray,
dynamic_source: mx.nd.NDArray,
source_length: mx.nd.NDArray,
previous_word_id: mx.nd.NDArray,
previous_hidden: mx.nd.NDArray,
decoder_states: List[mx.nd.NDArray],
bucket_key: int) -> Tuple[mx.nd.NDArray, mx.nd.NDArray,
mx.nd.NDArray, mx.nd.NDArray,
List[mx.nd.NDArray]]:
"""
Runs forward pass of the single-step decoder.
:param encoded_source: Encoded source sentence.
:param dynamic_source: Dynamic encoding of source sentence.
:param source_length: Source length.
:param previous_word_id: Previous predicted word id.
:param previous_hidden: Previous hidden decoder state.
:param decoder_states: Decoder states.
:param bucket_key: Bucket key.
:return: Probability distribution over next word, attention scores, dynamic source encoding,
next hidden state, next decoder states.
"""
data = [encoded_source,
dynamic_source,
source_length,
previous_word_id.as_in_context(self.context),
previous_hidden] + decoder_states
decoder_batch = mx.io.DataBatch(
data=data,
label=None, bucket_key=bucket_key, provide_data=self._get_decoder_data_shapes(bucket_key))
# run forward pass
self.decoder_module.forward(data_batch=decoder_batch, is_train=False)
# collect outputs
softmax_out, attention_probs, dynamic_source, next_hidden, *next_layer_states = \
self.decoder_module.get_outputs()
return softmax_out, attention_probs, dynamic_source, next_hidden, next_layer_states