code: def _build_keras_model(hidden_units, learning_rate):
      """Creates a DNN Keras model for classifying taxi data.
  
      Args:
        hidden_units: [int], the layer sizes of the DNN (input layer first).
        learning_rate: [float], learning rate of the Adam optimizer.
  
      Returns:
        A keras Model.
      """
      real_valued_columns = [
          tf.feature_column.numeric_column(key, shape=())
          for key in features.transformed_names(
              features.DENSE_FLOAT_FEATURE_KEYS
          )
      ]
      categorical_columns = [
          tf.feature_column.categorical_column_with_identity(  # pylint: disable=g-complex-comprehension
              key,
              num_buckets=features.VOCAB_SIZE + features.OOV_SIZE,
              default_value=0,
          )
          for key in features.transformed_names(
              features.VOCAB_FEATURE_KEYS
          )
      ]
      categorical_columns += [
          tf.feature_column.categorical_column_with_identity(  # pylint: disable=g-complex-comprehension
              key, num_buckets=num_buckets, default_value=0
          )
          for key, num_buckets in zip(
              features.transformed_names(features.BUCKET_FEATURE_KEYS),
              features.BUCKET_FEATURE_BUCKET_COUNT,
          )
      ]
      categorical_columns += [
          tf.feature_column.categorical_column_with_identity(  # pylint: disable=g-complex-comprehension
              key, num_buckets=num_buckets, default_value=0
          )
          for key, num_buckets in zip(
              features.transformed_names(
                  features.CATEGORICAL_FEATURE_KEYS
              ),
              features.CATEGORICAL_FEATURE_MAX_VALUES,
          )
      ]
      indicator_column = [
          tf.feature_column.indicator_column(categorical_column)
          for categorical_column in categorical_columns
      ]
  
      model = _wide_and_deep_classifier(
          # TODO(b/140320729) Replace with premade wide_and_deep keras model
          wide_columns=indicator_column,
          deep_columns=real_valued_columns,
          dnn_hidden_units=hidden_units,
          learning_rate=learning_rate,
      )
      return model
  
file_path: tfx/tfx/experimental/templates/taxi/models/keras_model/model.py
imports: from absl import logging
  import tensorflow
  import tensorflow_transform
  from tfx.experimental.templates.taxi.models import features
  from tfx.experimental.templates.taxi.models.keras_model import constants
  from tfx_bsl.public import tfxio
label_desc: Use image node_id 4 for calling a global function with example usage: _build_keras_model(hidden_units, learning_rate) and returns: model