Project import generated by Copybara.

PiperOrigin-RevId: 189351689

Author: tf-transform-team

RELEASE.md

@@ -4,13 +4,12 @@
 
 ## Bug Fixes and Other Changes
 * Trim min/max value in `tft.bucketize where the computed number of bucket
-  boundaries is more than requested.  Updated documentation to
-  clearly indicate that the number of buckets is computed using approximate
-  algorithms, and that computed number can be more or less than requested.
+  boundaries is more than requested. Updated documentation to clearly indicate
+  that the number of buckets is computed using approximate algorithms, and that
+  computed number can be more or less than requested.
 * Change the namespace used for Beam metrics from `tensorflow_transform` to
   `tfx.Transform`.
-* Export all required names from top level module, so only one import is needed
-  in user code.
+* Update Beam metrics to also log vocabulary sizes.
 
 ## Breaking changes
 

examples/census_example.py

@@ -32,6 +32,11 @@
 from tensorflow.contrib import lookup
 from tensorflow.contrib.learn.python.learn.utils import input_fn_utils
 
+from tensorflow_transform.beam import impl as beam_impl
+from tensorflow_transform.beam.tft_beam_io import transform_fn_io
+from tensorflow_transform.coders import csv_coder
+from tensorflow_transform.coders import example_proto_coder
+from tensorflow_transform.saved import saved_transform_io
 from tensorflow_transform.tf_metadata import dataset_metadata
 from tensorflow_transform.tf_metadata import dataset_schema
 from tensorflow_transform.tf_metadata import metadata_io
@@ -128,7 +133,7 @@ def convert_label(label):
   # The "with" block will create a pipeline, and run that pipeline at the exit
   # of the block.
   with beam.Pipeline() as pipeline:
-    with tft.Context(temp_dir=tempfile.mkdtemp()):
+    with beam_impl.Context(temp_dir=tempfile.mkdtemp()):
       # Create a coder to read the census data with the schema.  To do this we
       # need to list all columns in order since the schema doesn't specify the
       # order of columns in the csv.
@@ -138,7 +143,7 @@ def convert_label(label):
           'capital-gain', 'capital-loss', 'hours-per-week', 'native-country',
           'label'
       ]
-      converter = tft.CsvCoder(ordered_columns, RAW_DATA_METADATA.schema)
+      converter = csv_coder.CsvCoder(ordered_columns, RAW_DATA_METADATA.schema)
 
       # Read in raw data and convert using CSV converter.  Note that we apply
       # some Beam transformations here, which will not be encoded in the TF
@@ -159,12 +164,13 @@ def convert_label(label):
       # raw_data.
       raw_dataset = (raw_data, RAW_DATA_METADATA)
       transformed_dataset, transform_fn = (
-          raw_dataset | tft.AnalyzeAndTransformDataset(preprocessing_fn))
+          raw_dataset | beam_impl.AnalyzeAndTransformDataset(preprocessing_fn))
       transformed_data, transformed_metadata = transformed_dataset
 
       _ = transformed_data | 'WriteTrainData' >> tfrecordio.WriteToTFRecord(
           os.path.join(working_dir, TRANSFORMED_TRAIN_DATA_FILEBASE),
-          coder=tft.ExampleProtoCoder(transformed_metadata.schema))
+          coder=example_proto_coder.ExampleProtoCoder(
+              transformed_metadata.schema))
 
       # Now apply transform function to test data.  In this case we also remove
       # the header line from the CSV file and the trailing period at the end of
@@ -182,21 +188,22 @@ def convert_label(label):
       raw_test_dataset = (raw_test_data, RAW_DATA_METADATA)
 
       transformed_test_dataset = (
-          (raw_test_dataset, transform_fn) | tft.TransformDataset())
+          (raw_test_dataset, transform_fn) | beam_impl.TransformDataset())
       # Don't need transformed data schema, it's the same as before.
       transformed_test_data, _ = transformed_test_dataset
 
       _ = transformed_test_data | 'WriteTestData' >> tfrecordio.WriteToTFRecord(
           os.path.join(working_dir, TRANSFORMED_TEST_DATA_FILEBASE),
-          coder=tft.ExampleProtoCoder(transformed_metadata.schema))
+          coder=example_proto_coder.ExampleProtoCoder(
+              transformed_metadata.schema))
 
       # Will write a SavedModel and metadata to two subdirectories of
-      # working_dir, given by tft.TRANSFORM_FN_DIR and
-      # tft.TRANSFORMED_METADATA_DIR respectively.
+      # working_dir, given by transform_fn_io.TRANSFORM_FN_DIR and
+      # transform_fn_io.TRANSFORMED_METADATA_DIR respectively.
       _ = (
           transform_fn
           | 'WriteTransformFn' >>
-          tft.WriteTransformFn(working_dir))
+          transform_fn_io.WriteTransformFn(working_dir))
 
 # Functions for training
 
@@ -214,7 +221,8 @@ def _make_training_input_fn(working_dir, filebase, batch_size):
     The input function for training or eval.
   """
   transformed_metadata = metadata_io.read_metadata(
-      os.path.join(working_dir, tft.TRANSFORMED_METADATA_DIR))
+      os.path.join(
+          working_dir, transform_fn_io.TRANSFORMED_METADATA_DIR))
   transformed_feature_spec = transformed_metadata.schema.as_feature_spec()
 
   def input_fn():
@@ -257,8 +265,9 @@ def serving_input_fn():
     # Apply the transform function that was used to generate the materialized
     # data.
     _, transformed_features = (
-        tft.partially_apply_saved_transform(
-            os.path.join(working_dir, tft.TRANSFORM_FN_DIR), raw_features))
+        saved_transform_io.partially_apply_saved_transform(
+            os.path.join(working_dir, transform_fn_io.TRANSFORM_FN_DIR),
+            raw_features))
 
     return input_fn_utils.InputFnOps(transformed_features, None, default_inputs)
 

examples/sentiment_example.py

@@ -30,6 +30,10 @@
 from apache_beam.io import tfrecordio
 from tensorflow.contrib import learn
 from tensorflow.contrib.learn.python.learn.utils import input_fn_utils
+from tensorflow_transform.beam import impl as beam_impl
+from tensorflow_transform.beam.tft_beam_io import transform_fn_io
+from tensorflow_transform.coders import example_proto_coder
+from tensorflow_transform.saved import saved_transform_io
 from tensorflow_transform.tf_metadata import dataset_metadata
 from tensorflow_transform.tf_metadata import dataset_schema
 from tensorflow_transform.tf_metadata import metadata_io
@@ -137,14 +141,16 @@ def read_and_shuffle_data(
             (train_neg_filepattern, train_pos_filepattern))
         | 'WriteTrainData' >> tfrecordio.WriteToTFRecord(
             os.path.join(working_dir, SHUFFLED_TRAIN_DATA_FILEBASE),
-            coder=tft.ExampleProtoCoder(RAW_DATA_METADATA.schema)))
+            coder=example_proto_coder.ExampleProtoCoder(
+                RAW_DATA_METADATA.schema)))
     _ = (
         pipeline
         | 'ReadAndShuffleTest' >> ReadAndShuffleData(
             (test_neg_filepattern, test_pos_filepattern))
         | 'WriteTestData' >> tfrecordio.WriteToTFRecord(
             os.path.join(working_dir, SHUFFLED_TEST_DATA_FILEBASE),
-            coder=tft.ExampleProtoCoder(RAW_DATA_METADATA.schema)))
+            coder=example_proto_coder.ExampleProtoCoder(
+                RAW_DATA_METADATA.schema)))
     # pylint: enable=no-value-for-parameter
 
 
@@ -161,20 +167,22 @@ def transform_data(working_dir):
   """
 
   with beam.Pipeline() as pipeline:
-    with tft.Context(temp_dir=tempfile.mkdtemp()):
+    with beam_impl.Context(temp_dir=tempfile.mkdtemp()):
       train_data = (
           pipeline |
           'ReadTrain' >> tfrecordio.ReadFromTFRecord(
               os.path.join(working_dir,
                            SHUFFLED_TRAIN_DATA_FILEBASE + '*'),
-              coder=tft.ExampleProtoCoder(RAW_DATA_METADATA.schema)))
+              coder=example_proto_coder.ExampleProtoCoder(
+                  RAW_DATA_METADATA.schema)))
 
       test_data = (
           pipeline |
           'ReadTest' >> tfrecordio.ReadFromTFRecord(
               os.path.join(working_dir,
                            SHUFFLED_TEST_DATA_FILEBASE + '*'),
-              coder=tft.ExampleProtoCoder(RAW_DATA_METADATA.schema)))
+              coder=example_proto_coder.ExampleProtoCoder(
+                  RAW_DATA_METADATA.schema)))
 
       def preprocessing_fn(inputs):
         """Preprocess input columns into transformed columns."""
@@ -193,34 +201,36 @@ def preprocessing_fn(inputs):
 
       (transformed_train_data, transformed_metadata), transform_fn = (
           (train_data, RAW_DATA_METADATA)
-          | 'AnalyzeAndTransform' >> tft.AnalyzeAndTransformDataset(
+          | 'AnalyzeAndTransform' >> beam_impl.AnalyzeAndTransformDataset(
               preprocessing_fn))
 
       transformed_test_data, _ = (
           ((test_data, RAW_DATA_METADATA), transform_fn)
-          | 'Transform' >> tft.TransformDataset())
+          | 'Transform' >> beam_impl.TransformDataset())
 
       _ = (
           transformed_train_data
           | 'WriteTrainData' >> tfrecordio.WriteToTFRecord(
               os.path.join(working_dir,
                            TRANSFORMED_TRAIN_DATA_FILEBASE),
-              coder=tft.ExampleProtoCoder(transformed_metadata.schema)))
+              coder=example_proto_coder.ExampleProtoCoder(
+                  transformed_metadata.schema)))
 
       _ = (
           transformed_test_data
           | 'WriteTestData' >> tfrecordio.WriteToTFRecord(
               os.path.join(working_dir,
                            TRANSFORMED_TEST_DATA_FILEBASE),
-              coder=tft.ExampleProtoCoder(transformed_metadata.schema)))
+              coder=example_proto_coder.ExampleProtoCoder(
+                  transformed_metadata.schema)))
 
       # Will write a SavedModel and metadata to two subdirectories of
-      # working_dir, given by tft.TRANSFORM_FN_DIR and
-      # tft.TRANSFORMED_METADATA_DIR respectively.
+      # working_dir, given by transform_fn_io.TRANSFORM_FN_DIR and
+      # transform_fn_io.TRANSFORMED_METADATA_DIR respectively.
       _ = (
           transform_fn
           | 'WriteTransformFn' >>
-          tft.WriteTransformFn(working_dir))
+          transform_fn_io.WriteTransformFn(working_dir))
 
 
 # Functions for training
@@ -239,7 +249,8 @@ def _make_training_input_fn(working_dir, filebase, batch_size):
     The input function for training or eval.
   """
   transformed_metadata = metadata_io.read_metadata(
-      os.path.join(working_dir, tft.TRANSFORMED_METADATA_DIR))
+      os.path.join(
+          working_dir, transform_fn_io.TRANSFORMED_METADATA_DIR))
   transformed_feature_spec = transformed_metadata.schema.as_feature_spec()
 
   def input_fn():
@@ -282,8 +293,9 @@ def serving_input_fn():
     # Apply the transform function that was used to generate the materialized
     # data.
     _, transformed_features = (
-        tft.partially_apply_saved_transform(
-            os.path.join(working_dir, tft.TRANSFORM_FN_DIR), raw_features))
+        saved_transform_io.partially_apply_saved_transform(
+            os.path.join(working_dir, transform_fn_io.TRANSFORM_FN_DIR),
+            raw_features))
 
     return input_fn_utils.InputFnOps(transformed_features, None, default_inputs)
 

examples/simple_example.py

@@ -23,6 +23,7 @@
 
 import tensorflow as tf
 import tensorflow_transform as tft
+import tensorflow_transform.beam.impl as beam_impl
 from tensorflow_transform.tf_metadata import dataset_metadata
 from tensorflow_transform.tf_metadata import dataset_schema
 
@@ -59,9 +60,9 @@ def preprocessing_fn(inputs):
           tf.float32, [], dataset_schema.FixedColumnRepresentation())
   }))
 
-  with tft.Context(temp_dir=tempfile.mkdtemp()):
+  with beam_impl.Context(temp_dir=tempfile.mkdtemp()):
     transformed_dataset, transform_fn = (  # pylint: disable=unused-variable
-        (raw_data, raw_data_metadata) | tft.AnalyzeAndTransformDataset(
+        (raw_data, raw_data_metadata) | beam_impl.AnalyzeAndTransformDataset(
             preprocessing_fn))
 
   transformed_data, transformed_metadata = transformed_dataset  # pylint: disable=unused-variable

getting_started.md

@@ -120,7 +120,7 @@ raw_data = [
 
 raw_data_metadata = ...
 transformed_dataset, transform_fn = (
-    (raw_data, raw_data_metadata) | tft.AnalyzeAndTransformDataset(
+    (raw_data, raw_data_metadata) | beam_impl.AnalyzeAndTransformDataset(
         preprocessing_fn))
 transformed_data, transformed_metadata = transformed_dataset
 ```
@@ -278,7 +278,7 @@ between reading the lines of the CSV file, and applying the converter that
 converts each CSV row to an instance in the in-memory format.
 
 ```
-converter = tft.CsvCoder(ordered_columns, raw_data_schema)
+converter = csv_coder.CsvCoder(ordered_columns, raw_data_schema)
 
 raw_data = (
     p
@@ -351,13 +351,13 @@ shards that are written.
 ```
 transformed_data | "WriteTrainData" >> tfrecordio.WriteToTFRecord(
     transformed_eval_data_base,
-    coder=tft.ExampleProtoCoder(transformed_metadata))
+    coder=example_proto_coder.ExampleProtoCoder(transformed_metadata))
 ```
 
 In addition to the training data, we also write out the metadata.
 
 ```
-transformed_metadata | 'WriteMetadata' >> tft.WriteMetadata(
+transformed_metadata | 'WriteMetadata' >> beam_metadata_io.WriteMetadata(
     transformed_metadata_file, pipeline=p)
 ```
 

tensorflow_transform/__init__.py

@@ -16,18 +16,6 @@
 # pylint: disable=wildcard-import
 from tensorflow_transform.analyzers import *
 from tensorflow_transform.api import apply_function
-from tensorflow_transform.beam.impl import AnalyzeAndTransformDataset
-from tensorflow_transform.beam.impl import AnalyzeDataset
-from tensorflow_transform.beam.impl import Context
-from tensorflow_transform.beam.impl import TransformDataset
-from tensorflow_transform.beam.tft_beam_io.transform_fn_io import ReadTransformFn
-from tensorflow_transform.beam.tft_beam_io.transform_fn_io import TRANSFORM_FN_DIR
-from tensorflow_transform.beam.tft_beam_io.transform_fn_io import TRANSFORMED_METADATA_DIR
-from tensorflow_transform.beam.tft_beam_io.transform_fn_io import WriteTransformFn
-from tensorflow_transform.coders.csv_coder import CsvCoder
-from tensorflow_transform.coders.example_proto_coder import ExampleProtoCoder
 from tensorflow_transform.mappers import *
 from tensorflow_transform.pretrained_models import *
-from tensorflow_transform.saved.saved_transform_io import apply_saved_transform
-from tensorflow_transform.saved.saved_transform_io import partially_apply_saved_transform
 # pylint: enable=wildcard-import

tensorflow_transform/analyzers.py

@@ -28,6 +28,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import re
 
 import numpy as np
@@ -38,7 +39,19 @@
 VOCAB_FILENAME_PREFIX = 'vocab_'
 VOCAB_FREQUENCY_FILENAME_PREFIX = 'vocab_frequency_'
 
+# Named tuple with details for each output of an Analyzer.
+_AnalyzerOutputInfo = collections.namedtuple(
+    'AnalyzerOutputInfo', ['name', 'dtype', 'is_asset'])
 
+
+# NOTE: this code is designed so that Analyzer is pickleable, and in particular
+# does not try to pickle a tf.Graph or tf.Tensor which may not be pickleable.
+# This is due to https://issues.apache.org/jira/browse/BEAM-3812.  Until that
+# issue is fixed, anything that is a member variable of a Beam PTransform may
+# end up getting pickled.  Instances of Analyzer do end up as member variables
+# of a PTransform in our implementation of tf.Transform on Beam currently, so
+# we must avoid directly putting `Tensor`s inside `Analyzer`, and instead use
+# tensor names.
 class Analyzer(object):
   """An operation-like class for full-pass analyses of data.
 
@@ -68,28 +81,37 @@ def __init__(self, inputs, output_dtype_shape_and_is_asset, spec, name):
     for tensor in inputs:
       if not isinstance(tensor, tf.Tensor):
         raise ValueError('Analyzers can only accept `Tensor`s as inputs')
-    self._inputs = inputs
-    self._outputs = []
-    self._output_is_asset_map = {}
+    self._input_tensor_names = [tensor.name for tensor in inputs]
+    self._output_infos = []
     with tf.name_scope(name) as scope:
       self._name = scope
       for dtype, shape, is_asset in output_dtype_shape_and_is_asset:
         output_tensor = tf.placeholder(dtype, shape)
         if is_asset and output_tensor.dtype != tf.string:
           raise ValueError(('Tensor {} cannot represent an asset, because it '
                             'is not a string.').format(output_tensor.name))
-        self._outputs.append(output_tensor)
-        self._output_is_asset_map[output_tensor] = is_asset
+        self._output_infos.append(_AnalyzerOutputInfo(
+            output_tensor.name, output_tensor.dtype, is_asset))
     self._spec = spec
     tf.add_to_collection(ANALYZER_COLLECTION, self)
 
+  @property
+  def input_tensor_names(self):
+    return self._input_tensor_names
+
+  @property
+  def output_infos(self):
+    return self._output_infos
+
   @property
   def inputs(self):
-    return self._inputs
+    return [tf.get_default_graph().get_tensor_by_name(name)
+            for name in self._input_tensor_names]
 
   @property
   def outputs(self):
-    return self._outputs
+    return [tf.get_default_graph().get_tensor_by_name(output_info.name)
+            for output_info in self._output_infos]
 
   @property
   def spec(self):
@@ -99,9 +121,6 @@ def spec(self):
   def name(self):
     return self._name
 
-  def output_is_asset(self, output_tensor):
-    return self._output_is_asset_map[output_tensor]
-
 
 class CombinerSpec(object):
   """Analyze using combiner function.