def test_transformer_decoder(): reset_session() # Construct the layer from rinokeras.core.v1x.models.transformer import TransformerDecoder, TransformerInputEmbedding tie = TransformerInputEmbedding(128, False) layer = TransformerDecoder(embedding_layer=tie, output_layer=tf.keras.layers.Dense(128), n_layers=2, n_heads=4, d_model=128, d_filter=32) # Make sure that the layer is not None assert layer is not None # Encoded values source_tensor, _ = random_tensor((2,32,128)) target_tensor, _ = random_tensor((2,32,64)) # Get random masking values source_mask, _ = random_mask_tensor(2, 32) target_mask, _ = random_mask_tensor(2, 32) source_mask = convert_to_attention_mask(source_tensor, source_mask) target_mask = convert_to_attention_mask(target_tensor, target_mask) # Get the output of the layer value = layer((source_tensor, target_tensor), mask=(source_mask, target_mask)) # Create a named temporary file for save/restore testing with tempfile.TemporaryFile() as weights_file: # Construct the session output = run_simple_session_save_weights(inputs=[value], feed={}, weights=[layer], weights_file=weights_file) assert_not_none(output) assert_expected_shapes(output, [(2,32,128)]) # Check loading and restoring load_restore_test(output=output, inputs=[value], feed={}, weights=[layer], weights_file=weights_file) # Do regression testing check_regression('transformer_decoder_output', output, __file__, 'regression_outputs/test_transformer_decoder_outputs.json', debug=_RK_REBUILD_REGRESSION) # Do a config test from_config_test(TransformerDecoder, layer)
def test_transformer_encoder_masking_with_conv(): reset_session() # Construct the layer from rinokeras.core.v1x.models.transformer import TransformerEncoder layer = TransformerEncoder(embedding_layer=tf.keras.layers.Dense(64), n_layers=2, n_heads=4, d_model=64, d_filter=128) # Make sure that the layer is not None assert layer is not None # Encoded values input_tensor, _ = random_tensor((16, 32, 128)) input_mask, _ = random_mask_tensor(16, 32) input_mask = convert_to_attention_mask(input_tensor, input_mask) conv_mask, _ = random_mask_tensor(16, 32) conv_mask = convert_sequence_length_to_sequence_mask( input_tensor, conv_mask) # Get the output of the layer value = layer(input_tensor, mask=(input_mask, conv_mask)) # Create a named temporary file for save/restore testing with tempfile.TemporaryFile() as weights_file: # Construct the session output = run_simple_session_save_weights(inputs=[value], feed={}, weights=[layer], weights_file=weights_file) assert_not_none(output) assert_expected_shapes(output, [(16, 32, 64)]) # Check loading and restoring load_restore_test(output=output, inputs=[value], feed={}, weights=[layer], weights_file=weights_file) # Do regression testing check_regression( 'transformer_encoder_output_masking_with_conv', output, __file__, 'regression_outputs/test_transformer_encoder_outputs.json', debug=_RK_REBUILD_REGRESSION) # Do a config test from_config_test(TransformerEncoder, layer)
def test_convert_to_attention_mask_2(): reset_session() # Construct the layer from rinokeras.core.v1x.utils import convert_to_attention_mask # Encoded values input_tensor, _ = random_sequence_tensor(2, 32, 128) input_mask, _ = random_mask_tensor(2, 32) # Get the output of the layer value = convert_to_attention_mask(input_tensor, input_mask) # Construct the session output = run_simple_session(inputs=[value], feed={}) assert_not_none(output) assert_expected_shapes(output, [(2, 32, 32)]) # Do regression testing check_regression('convert_to_attention_mask_2', output, __file__, 'regression_outputs/test_utils_outputs.json', debug=_RK_REBUILD_REGRESSION)