Spaces:
Running
Running
| # Copyright 2020 The TensorFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Tests for nlp.nhnet.decoder.""" | |
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| import numpy as np | |
| import tensorflow as tf | |
| from official.nlp.modeling import layers | |
| from official.nlp.nhnet import configs | |
| from official.nlp.nhnet import decoder | |
| from official.nlp.nhnet import utils | |
| class DecoderTest(tf.test.TestCase): | |
| def setUp(self): | |
| super(DecoderTest, self).setUp() | |
| self._config = utils.get_test_params() | |
| def test_transformer_decoder(self): | |
| decoder_block = decoder.TransformerDecoder( | |
| num_hidden_layers=self._config.num_hidden_layers, | |
| hidden_size=self._config.hidden_size, | |
| num_attention_heads=self._config.num_attention_heads, | |
| intermediate_size=self._config.intermediate_size, | |
| intermediate_activation=self._config.hidden_act, | |
| hidden_dropout_prob=self._config.hidden_dropout_prob, | |
| attention_probs_dropout_prob=self._config.attention_probs_dropout_prob, | |
| initializer_range=self._config.initializer_range) | |
| decoder_block.build(None) | |
| self.assertEqual(len(decoder_block.layers), self._config.num_hidden_layers) | |
| def test_bert_decoder(self): | |
| seq_length = 10 | |
| encoder_input_ids = tf.keras.layers.Input( | |
| shape=(seq_length,), name="encoder_input_ids", dtype=tf.int32) | |
| target_ids = tf.keras.layers.Input( | |
| shape=(seq_length,), name="target_ids", dtype=tf.int32) | |
| encoder_outputs = tf.keras.layers.Input( | |
| shape=(seq_length, self._config.hidden_size), | |
| name="all_encoder_outputs", | |
| dtype=tf.float32) | |
| embedding_lookup = layers.OnDeviceEmbedding( | |
| vocab_size=self._config.vocab_size, | |
| embedding_width=self._config.hidden_size, | |
| initializer=tf.keras.initializers.TruncatedNormal( | |
| stddev=self._config.initializer_range), | |
| name="word_embeddings") | |
| cross_attention_bias = decoder.AttentionBias(bias_type="single_cross")( | |
| encoder_input_ids) | |
| self_attention_bias = decoder.AttentionBias(bias_type="decoder_self")( | |
| target_ids) | |
| inputs = dict( | |
| attention_bias=cross_attention_bias, | |
| self_attention_bias=self_attention_bias, | |
| target_ids=target_ids, | |
| all_encoder_outputs=encoder_outputs) | |
| decoder_layer = decoder.Decoder(self._config, embedding_lookup) | |
| outputs = decoder_layer(inputs) | |
| model_inputs = dict( | |
| encoder_input_ids=encoder_input_ids, | |
| target_ids=target_ids, | |
| all_encoder_outputs=encoder_outputs) | |
| model = tf.keras.Model(inputs=model_inputs, outputs=outputs, name="test") | |
| self.assertLen(decoder_layer.trainable_weights, 30) | |
| # Forward path. | |
| fake_inputs = { | |
| "encoder_input_ids": np.zeros((2, 10), dtype=np.int32), | |
| "target_ids": np.zeros((2, 10), dtype=np.int32), | |
| "all_encoder_outputs": np.zeros((2, 10, 16), dtype=np.float32), | |
| } | |
| output_tensor = model(fake_inputs) | |
| self.assertEqual(output_tensor.shape, (2, 10, 16)) | |
| def test_multi_doc_decoder(self): | |
| self._config = utils.get_test_params(cls=configs.NHNetConfig) | |
| seq_length = 10 | |
| num_docs = 5 | |
| encoder_input_ids = tf.keras.layers.Input( | |
| shape=(num_docs, seq_length), name="encoder_input_ids", dtype=tf.int32) | |
| target_ids = tf.keras.layers.Input( | |
| shape=(seq_length,), name="target_ids", dtype=tf.int32) | |
| encoder_outputs = tf.keras.layers.Input( | |
| shape=(num_docs, seq_length, self._config.hidden_size), | |
| name="all_encoder_outputs", | |
| dtype=tf.float32) | |
| embedding_lookup = layers.OnDeviceEmbedding( | |
| vocab_size=self._config.vocab_size, | |
| embedding_width=self._config.hidden_size, | |
| initializer=tf.keras.initializers.TruncatedNormal( | |
| stddev=self._config.initializer_range), | |
| name="word_embeddings") | |
| doc_attention_probs = tf.keras.layers.Input( | |
| shape=(self._config.num_decoder_attn_heads, seq_length, num_docs), | |
| name="doc_attention_probs", | |
| dtype=tf.float32) | |
| cross_attention_bias = decoder.AttentionBias(bias_type="multi_cross")( | |
| encoder_input_ids) | |
| self_attention_bias = decoder.AttentionBias(bias_type="decoder_self")( | |
| target_ids) | |
| inputs = dict( | |
| attention_bias=cross_attention_bias, | |
| self_attention_bias=self_attention_bias, | |
| target_ids=target_ids, | |
| all_encoder_outputs=encoder_outputs, | |
| doc_attention_probs=doc_attention_probs) | |
| decoder_layer = decoder.Decoder(self._config, embedding_lookup) | |
| outputs = decoder_layer(inputs) | |
| model_inputs = dict( | |
| encoder_input_ids=encoder_input_ids, | |
| target_ids=target_ids, | |
| all_encoder_outputs=encoder_outputs, | |
| doc_attention_probs=doc_attention_probs) | |
| model = tf.keras.Model(inputs=model_inputs, outputs=outputs, name="test") | |
| self.assertLen(decoder_layer.trainable_weights, 30) | |
| # Forward path. | |
| fake_inputs = { | |
| "encoder_input_ids": | |
| np.zeros((2, num_docs, seq_length), dtype=np.int32), | |
| "target_ids": | |
| np.zeros((2, seq_length), dtype=np.int32), | |
| "all_encoder_outputs": | |
| np.zeros((2, num_docs, seq_length, 16), dtype=np.float32), | |
| "doc_attention_probs": | |
| np.zeros( | |
| (2, self._config.num_decoder_attn_heads, seq_length, num_docs), | |
| dtype=np.float32) | |
| } | |
| output_tensor = model(fake_inputs) | |
| self.assertEqual(output_tensor.shape, (2, seq_length, 16)) | |
| if __name__ == "__main__": | |
| tf.test.main() | |