Fix list comprehension for bucket features in tfx template.

PiperOrigin-RevId: 312005513
m2kar · May 18, 2020 · 62ba786 · 62ba786
1 parent 78e4c70
commit 62ba786
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/tfx/experimental/templates/taxi/models/features.py b/tfx/experimental/templates/taxi/models/features.py
@@ -32,10 +32,10 @@
 
 # Name of features which have continuous float values. These features will be
 # bucketized using `tft.bucketize`, and will be used as categorical features.
-BUCKET_FEATURE_KEYS = []
+BUCKET_FEATURE_KEYS = ['pickup_latitude']
 # Number of buckets used by tf.transform for encoding each feature. The length
 # of this list should be the same with BUCKET_FEATURE_KEYS.
-BUCKET_FEATURE_BUCKET_COUNT = []
+BUCKET_FEATURE_BUCKET_COUNT = [10]
 
 # Name of features which have categorical values which are mapped to integers.
 # These features will be used as categorical features.

diff --git a/tfx/experimental/templates/taxi/models/keras/model.py b/tfx/experimental/templates/taxi/models/keras/model.py
@@ -106,9 +106,10 @@ def _build_keras_model(hidden_units, learning_rate):
   categorical_columns += [
       tf.feature_column.categorical_column_with_identity(  # pylint: disable=g-complex-comprehension
           key,
-          num_buckets=features.BUCKET_FEATURE_BUCKET_COUNT,
-          default_value=0)
-      for key in features.transformed_names(features.BUCKET_FEATURE_KEYS)
+          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

diff --git a/tfx/experimental/templates/taxi/models/keras/model_test.py b/tfx/experimental/templates/taxi/models/keras/model_test.py
@@ -27,10 +27,10 @@ class ModelTest(tf.test.TestCase):
   def testBuildKerasModel(self):
     built_model = model._build_keras_model(
         hidden_units=[1, 1], learning_rate=0.1)  # pylint: disable=protected-access
-    self.assertEqual(len(built_model.layers), 8)
+    self.assertEqual(len(built_model.layers), 9)
 
     built_model = model._build_keras_model(hidden_units=[1], learning_rate=0.1)  # pylint: disable=protected-access
-    self.assertEqual(len(built_model.layers), 7)
+    self.assertEqual(len(built_model.layers), 8)
 
 
 if __name__ == '__main__':