# Path: tensorflow_transform/saved/saved_transform_io.py
# _MANGLED_TENSOR_NAME_RE = re.compile(
#     r'(.*)\$(indices|values|dense_shape|dense_tensor)$')
# _PARTITIONED_VARIABLE_NAME_RE = re.compile(r'^(.*)/part_(\d*)$')
# def _update_legacy_signature(signature):
# def _load_transform_saved_model(transform_savedmodel_dir):
# def _expand_input_map(logical_input_map, input_signature):
# def _maybe_register_addon_ops():
#   def _try_import(name):
# def _partially_apply_saved_transform_impl(saved_model_dir,
#                                           logical_input_map,
#                                           tensor_replacement_map=None):
#   def lookup_remapped_tensor(tensor_name):
#   def lookup_tensor_or_sparse_or_composite_tensor(tensor_info):
# def partially_apply_saved_transform_internal(saved_model_dir,
#                                              logical_input_map,
#                                              tensor_replacement_map=None):
# def write_saved_transform_from_session(
#     session, inputs, outputs, export_path, as_text=False):
# def exported_as_v1(transform_savedmodel_dir):


import os
import tempfile
import numpy as np
import tensorflow as tf
from tensorflow_transform.saved import saved_transform_io
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import lookup_ops
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for saved_transform_io."""



# pylint: disable=g-direct-tensorflow-import
# pylint: enable=g-direct-tensorflow-import


# TODO(b/123241798): Find an open-source compatible way to access
# FLAGS.test_tmpdir.
def _create_test_saved_model():
  export_path = os.path.join(tempfile.mkdtemp(), 'export')

  with tf.compat.v1.Graph().as_default():
    with tf.compat.v1.Session().as_default() as session:
      input_float = tf.compat.v1.placeholder(tf.float32, shape=[1])
      output = (input_float - 2.0) / 5.0
      inputs = {'x': input_float}
      outputs = {'x_scaled': output}

# Path: tensorflow_transform/saved/saved_transform_io.py
# _MANGLED_TENSOR_NAME_RE = re.compile(
#     r'(.*)\$(indices|values|dense_shape|dense_tensor)$')
# _PARTITIONED_VARIABLE_NAME_RE = re.compile(r'^(.*)/part_(\d*)$')
# def _update_legacy_signature(signature):
# def _load_transform_saved_model(transform_savedmodel_dir):
# def _expand_input_map(logical_input_map, input_signature):
# def _maybe_register_addon_ops():
#   def _try_import(name):
# def _partially_apply_saved_transform_impl(saved_model_dir,
#                                           logical_input_map,
#                                           tensor_replacement_map=None):
#   def lookup_remapped_tensor(tensor_name):
#   def lookup_tensor_or_sparse_or_composite_tensor(tensor_info):
# def partially_apply_saved_transform_internal(saved_model_dir,
#                                              logical_input_map,
#                                              tensor_replacement_map=None):
# def write_saved_transform_from_session(
#     session, inputs, outputs, export_path, as_text=False):
# def exported_as_v1(transform_savedmodel_dir):


# Path: tensorflow_transform/saved/saved_transform_io_test.py
import os
import tempfile
import numpy as np
import tensorflow as tf
from tensorflow_transform.saved import saved_transform_io
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import lookup_ops

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for saved_transform_io."""



# pylint: disable=g-direct-tensorflow-import
# pylint: enable=g-direct-tensorflow-import


# TODO(b/123241798): Find an open-source compatible way to access
# FLAGS.test_tmpdir.
def _create_test_saved_model():
  export_path = os.path.join(tempfile.mkdtemp(), 'export')

  with tf.compat.v1.Graph().as_default():
    with tf.compat.v1.Session().as_default() as session:
      input_float = tf.compat.v1.placeholder(tf.float32, shape=[1])
      output = (input_float - 2.0) / 5.0
      inputs = {'x': input_float}
      outputs = {'x_scaled': output}

      saved_transform_io.write_saved_transform_from_session(