def test_distribution_strategy_output_with_adapt(self, strategy): # TODO(b/180614455): remove this check when MLIR bridge is always enabled. if backend.is_tpu_strategy(strategy): self.skipTest("This test needs MLIR bridge on TPU.") vocab_data = [[ "earth", "earth", "earth", "earth", "wind", "wind", "wind", "and", "and", "fire" ]] vocab_dataset = tf.data.Dataset.from_tensors(vocab_data) input_array = np.array([["earth", "wind", "and", "fire"], ["fire", "and", "earth", "michigan"]]) input_dataset = tf.data.Dataset.from_tensor_slices(input_array).batch( 2, drop_remainder=True) expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]] tf.config.set_soft_device_placement(True) with strategy.scope(): input_data = keras.Input(shape=(None, ), dtype=tf.string) layer = text_vectorization.TextVectorization( max_tokens=None, standardize=None, split=None, output_mode=text_vectorization.INT) layer.adapt(vocab_dataset) int_data = layer(input_data) model = keras.Model(inputs=input_data, outputs=int_data) output_dataset = model.predict(input_dataset) self.assertAllEqual(expected_output, output_dataset)
def test_distribution_strategy_output_with_adapt(self, distribution): vocab_data = [[ "earth", "earth", "earth", "earth", "wind", "wind", "wind", "and", "and", "fire" ]] vocab_dataset = tf.data.Dataset.from_tensors(vocab_data) input_array = np.array([["earth", "wind", "and", "fire"], ["fire", "and", "earth", "michigan"]]) input_dataset = tf.data.Dataset.from_tensor_slices(input_array).batch( 2, drop_remainder=True) expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]] tf.config.set_soft_device_placement(True) with distribution.scope(): input_data = keras.Input(shape=(None,), dtype=tf.string) layer = text_vectorization.TextVectorization( max_tokens=None, standardize=None, split=None, output_mode=text_vectorization.INT) layer.adapt(vocab_dataset) int_data = layer(input_data) model = keras.Model(inputs=input_data, outputs=int_data) output_dataset = model.predict(input_dataset) self.assertAllEqual(expected_output, output_dataset)
def test_distribution_strategy_output(self, strategy): if (backend.is_tpu_strategy(strategy) and not tf_test_utils.is_mlir_bridge_enabled()): self.skipTest("TPU tests require MLIR bridge") vocab_data = ["earth", "wind", "and", "fire"] input_array = np.array([["earth", "wind", "and", "fire"], ["fire", "and", "earth", "michigan"]]) input_dataset = tf.data.Dataset.from_tensor_slices(input_array).batch( 2, drop_remainder=True) expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]] tf.config.set_soft_device_placement(True) with strategy.scope(): input_data = keras.Input(shape=(None, ), dtype=tf.string) layer = text_vectorization.TextVectorization( max_tokens=None, standardize=None, split=None, output_mode=text_vectorization.INT, vocabulary=vocab_data) int_data = layer(input_data) model = keras.Model(inputs=input_data, outputs=int_data) output_dataset = model.predict(input_dataset) self.assertAllEqual(expected_output, output_dataset)