message: Update the tutorial to simplify the datasets creation.
diff: @@ -66,6 +66,7 @@ def make_coordinator(num_workers, num_ps):
  parameter_server_strategy_v2.ParameterServerStrategyV2(cluster_resolver))
  +# TODO(yuefengz): move this to keras/integration_tests.
  class KPLTest(test.TestCase, parameterized.TestCase):
  @classmethod
  @@ -98,7 +99,7 @@ class KPLTest(test.TestCase, parameterized.TestCase):
  feature_ps = keras.Model({"features": raw_feature_input}, feature_id_input)
  raw_label_input = keras.layers.Input(
  - shape=(), dtype=dtypes.string, name="label")
  + shape=(1,), dtype=dtypes.string, name="label")
  label_id_input = label_lookup_layer(raw_label_input)
  label_ps = keras.Model({"label": raw_label_input}, label_id_input)
  @@ -123,30 +124,23 @@ class KPLTest(test.TestCase, parameterized.TestCase):
  def feature_and_label_gen():
  while True:
  features = random.sample(FEATURE_VOCAB, 3)
  - label = "yes" if "avenger" in features else "no"
  + label = ["yes"] if "avenger" in features else ["no"]
  yield {"features": features, "label": label}
  - # The dataset will be created on the coordinator?
  + # The dataset will be created on the coordinator.
  raw_dataset = dataset_ops.Dataset.from_generator(
  feature_and_label_gen,
  - output_types={
  - "features": dtypes.string,
  - "label": dtypes.string
  - }).shuffle(200).batch(32)
  - preproc_dataset = raw_dataset.map(
  - lambda x: { # pylint: disable=g-long-lambda
  - "features": feature_ps(x["features"]),
  - "label": label_ps(x["label"])
  - })
  - train_dataset = preproc_dataset.map(lambda x: ( # pylint: disable=g-long-lambda
  + output_signature={
  + "features": tensor_spec.TensorSpec([3], dtypes.string),
  + "label": tensor_spec.TensorSpec([1], dtypes.string)
  + }).shuffle(100).batch(32)
  +
  + train_dataset = raw_dataset.map(lambda x: ( # pylint: disable=g-long-lambda
  {
  - "features": x["features"]
  - }, [x["label"]]))
  + "features": feature_ps(x["features"])
  + }, label_ps(x["label"])))
  return train_dataset
  - distributed_dataset = self.coordinator.create_per_worker_dataset(
  - dataset_fn)
  -
  # Create the model. The input needs to be compatible with KPLs.
  model_input = keras.layers.Input(
  shape=(3,), dtype=dtypes.int64, name="model_input")
  @@ -161,7 +155,7 @@ class KPLTest(test.TestCase, parameterized.TestCase):
  emb_output)
  model = keras.Model({"features": model_input}, dense_output)
  - optimizer = rmsprop.RMSprop(learning_rate=0.01)
  + optimizer = rmsprop.RMSprop(learning_rate=0.1)
  accuracy = keras.metrics.Accuracy()
  @def_function.function
  @@ -183,11 +177,14 @@ class KPLTest(test.TestCase, parameterized.TestCase):
  self.coordinator._strategy.run(replica_fn, args=(iterator,))
  + distributed_dataset = self.coordinator.create_per_worker_dataset(dataset_fn)
  distributed_iterator = iter(distributed_dataset)
  - for _ in range(10):
  + for _ in range(4):
  + accuracy.reset_states()
  + for _ in range(7):
  self.coordinator.schedule(worker_fn, args=(distributed_iterator,))
  self.coordinator.join()
  - self.assertGreater(accuracy.result().numpy(), 0.0)
  + self.assertGreater(accuracy.result().numpy(), 0.5)
  # Create a saved model.
  model.feature_ps = feature_ps
  