Removing pretrained_models APIs from it since those are TF 1.x specific.

PiperOrigin-RevId: 367403376
tensorflow · Apr 8, 2021 · c7c9573 · c7c9573
1 parent 9d93810
commit c7c9573
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -13,6 +13,9 @@
 
 *   `DatasetMetadata` no longer accepts a dict as its input schema. `schema` is
     expected to be a `Schema` proto now.
+*   TF 1.15 specific APIs `apply_saved_model` and
+    `apply_function_with_checkpoint` were removed from the `tft` namespace.
+    They are still available under the `pretrained_models` module.
 
 ## Deprecations
 

diff --git a/tensorflow_transform/__init__.py b/tensorflow_transform/__init__.py
@@ -22,7 +22,6 @@
 from tensorflow_transform.mappers import *
 from tensorflow_transform.output_wrapper import TFTransformOutput
 from tensorflow_transform.output_wrapper import TransformFeaturesLayer
-from tensorflow_transform.pretrained_models import *
 from tensorflow_transform.py_func.api import apply_pyfunc
 # pylint: enable=wildcard-import
 

diff --git a/tensorflow_transform/beam/impl_test.py b/tensorflow_transform/beam/impl_test.py
@@ -32,6 +32,7 @@
 import tensorflow_transform as tft
 from tensorflow_transform import analyzers
 from tensorflow_transform import common
+from tensorflow_transform import pretrained_models
 from tensorflow_transform import schema_inference
 import tensorflow_transform.beam as tft_beam
 from tensorflow_transform.beam import tft_unit
@@ -180,7 +181,7 @@ def save_model_with_single_input(instance, export_dir):
 
     def preprocessing_fn(inputs):
       x = inputs['x']
-      output_col = tft.apply_saved_model(
+      output_col = pretrained_models.apply_saved_model(
           export_dir, x, tags=[tf.saved_model.SERVING])
       return {'out': output_col}
 
@@ -233,7 +234,7 @@ def save_model_with_hash_table(instance, export_dir):
 
     def preprocessing_fn(inputs):
       x = inputs['x']
-      output_col = tft.apply_saved_model(
+      output_col = pretrained_models.apply_saved_model(
           export_dir, x, tags=[tf.saved_model.SERVING])
       return {'out': output_col}
 
@@ -292,7 +293,7 @@ def preprocessing_fn(inputs):
       x = inputs['x']
       y = inputs['y']
       z = inputs['z']
-      sum_column = tft.apply_saved_model(
+      sum_column = pretrained_models.apply_saved_model(
           export_dir, {
               'name1': x,
               'name3': z,
@@ -350,7 +351,7 @@ def save_checkpoint(instance, checkpoint_path):
     def preprocessing_fn(inputs):
       x = inputs['x']
       y = inputs['y']
-      out_value = tft.apply_function_with_checkpoint(
+      out_value = pretrained_models.apply_function_with_checkpoint(
           tensor_fn, [x, y], checkpoint_path)
       return {'out': out_value}