|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Test decoding utility methods."""
|
|
|
| import abc
|
| import tensorflow as tf, tf_keras
|
|
|
| from official.nlp.modeling.ops import decoding_module
|
|
|
|
|
| def length_normalization(length, dtype):
|
| """Return length normalization factor."""
|
| return tf.pow(((5. + tf.cast(length, dtype)) / 6.), 0.0)
|
|
|
|
|
| class TestSubclass(decoding_module.DecodingModule, metaclass=abc.ABCMeta):
|
|
|
| def __init__(self,
|
| length_normalization_fn=length_normalization,
|
| extra_cache_output=True,
|
| dtype=tf.float32):
|
| super(TestSubclass, self).__init__(
|
| length_normalization_fn=length_normalization, dtype=dtype)
|
|
|
| def _create_initial_state(self, initial_ids, initial_cache, batch_size):
|
| pass
|
|
|
| def _grow_alive_seq(self, state, batch_size):
|
| pass
|
|
|
| def _process_finished_state(self, finished_state):
|
| pass
|
|
|
| def _get_new_finished_state(self, state, new_seq, new_log_probs,
|
| new_finished_flags, batch_size):
|
| pass
|
|
|
| def _finished_flags(self, topk_ids, state):
|
| pass
|
|
|
| def _continue_search(self, state):
|
| pass
|
|
|
| def _get_new_alive_state(self, new_seq, new_log_probs, new_finished_flags,
|
| new_cache):
|
| pass
|
|
|
|
|
| class DecodingModuleTest(tf.test.TestCase):
|
|
|
| def test_get_shape_keep_last_dim(self):
|
| y = tf.constant(4.0)
|
| x = tf.ones([7, tf.cast(tf.sqrt(y), tf.int32), 2, 5])
|
| shape = decoding_module.get_shape_keep_last_dim(x)
|
| self.assertAllEqual([None, None, None, 5], shape.as_list())
|
|
|
| def test_shape_list(self):
|
| x = tf.ones([7, 1])
|
| shape = decoding_module.shape_list(x)
|
| self.assertAllEqual([7, 1], shape)
|
|
|
| def test_inf(self):
|
| d = TestSubclass()
|
| inf_value = d.inf()
|
| self.assertAllEqual(inf_value, tf.constant(10000000., tf.float32))
|
|
|
| def test_length_normalization(self):
|
| d = TestSubclass()
|
| normalized_length = d.length_normalization_fn(32, tf.float32)
|
| self.assertAllEqual(normalized_length, tf.constant(1.0, tf.float32))
|
|
|
| if __name__ == '__main__':
|
| tf.test.main()
|
|
|
