# Path: tensorflow_transform/py_func/pyfunc_helper.py
# _PYFUNC_COLLECTION_KEY = 'pyfuncs'
# class _PyFuncDef(tfx_namedtuple.namedtuple('_PyFuncDef', ['token', 'func'])):
#   def from_proto(attr_value, import_scope=None):
#   def from_proto_string(proto_str, import_scope=None):
#   def to_proto(self, export_scope=None):
# def insert_pyfunc(func, Tout, stateful, name, *args):  # pylint: disable=invalid-name
# def register_pyfuncs_from_saved_transform(graph, meta_graph, loaded_in_tf2):
# 
# Path: tensorflow_transform/saved/constants.py
# TRANSFORM_TAG = 'transform'
# TRANSFORM_SIGNATURE = 'transform_signature'
# 
# Path: tensorflow_transform/saved/saved_model_loader.py
# def parse_saved_model(saved_model_dir):
# def _choose_meta_graph_def_internal(saved_model, tags):
# def choose_meta_graph_def(saved_model):
# def choose_meta_graph_def_and_raise(saved_model):
# def get_asset_tensors(saved_model_dir, meta_graph_def_to_load):


import importlib
import os
import re
import tensorflow as tf
from tensorflow_transform.py_func import pyfunc_helper
from tensorflow_transform.saved import constants
from tensorflow_transform.saved import saved_model_loader
from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.training import saver as tf_saver
      the corresponding Tensors must have the expected types and shapes.
    tensor_replacement_map: a dict of tensor names to `Tensors`.

  Returns:
    A pair of (unbound_inputs, outputs) where unbound_inputs is a dict of
    logical name to Tensors that are yet to be mapped or fed, and outputs is
    a dict of logical name to Tensor, as provided by the output signature
    of the transform graph

  Raises:
    ValueError: if the provided input_tensors dict has keys that are not part
      of the input signature, or any of the provided inputs have the wrong
      type or shape.
    RuntimeError: if there is no default graph available to which to apply the
      transform.
  """
  unbound_inputs, outputs = _partially_apply_saved_transform_impl(
      saved_model_dir, logical_input_map, tensor_replacement_map)
  return unbound_inputs, outputs


def write_saved_transform_from_session(
    session, inputs, outputs, export_path, as_text=False):
  """Write the current session as a SavedModel."""
  predict_signature_def = (
      tf.compat.v1.saved_model.signature_def_utils.predict_signature_def(
          inputs, outputs))

  builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_path)
  builder.add_meta_graph_and_variables(

# Path: tensorflow_transform/py_func/pyfunc_helper.py
# _PYFUNC_COLLECTION_KEY = 'pyfuncs'
# class _PyFuncDef(tfx_namedtuple.namedtuple('_PyFuncDef', ['token', 'func'])):
#   def from_proto(attr_value, import_scope=None):
#   def from_proto_string(proto_str, import_scope=None):
#   def to_proto(self, export_scope=None):
# def insert_pyfunc(func, Tout, stateful, name, *args):  # pylint: disable=invalid-name
# def register_pyfuncs_from_saved_transform(graph, meta_graph, loaded_in_tf2):
# 
# Path: tensorflow_transform/saved/constants.py
# TRANSFORM_TAG = 'transform'
# TRANSFORM_SIGNATURE = 'transform_signature'
# 
# Path: tensorflow_transform/saved/saved_model_loader.py
# def parse_saved_model(saved_model_dir):
# def _choose_meta_graph_def_internal(saved_model, tags):
# def choose_meta_graph_def(saved_model):
# def choose_meta_graph_def_and_raise(saved_model):
# def get_asset_tensors(saved_model_dir, meta_graph_def_to_load):


# Path: tensorflow_transform/saved/saved_transform_io.py
import importlib
import os
import re
import tensorflow as tf
from tensorflow_transform.py_func import pyfunc_helper
from tensorflow_transform.saved import constants
from tensorflow_transform.saved import saved_model_loader
from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.training import saver as tf_saver

      the corresponding Tensors must have the expected types and shapes.
    tensor_replacement_map: a dict of tensor names to `Tensors`.

  Returns:
    A pair of (unbound_inputs, outputs) where unbound_inputs is a dict of
    logical name to Tensors that are yet to be mapped or fed, and outputs is
    a dict of logical name to Tensor, as provided by the output signature
    of the transform graph

  Raises:
    ValueError: if the provided input_tensors dict has keys that are not part
      of the input signature, or any of the provided inputs have the wrong
      type or shape.
    RuntimeError: if there is no default graph available to which to apply the
      transform.
  """
  unbound_inputs, outputs = _partially_apply_saved_transform_impl(
      saved_model_dir, logical_input_map, tensor_replacement_map)
  return unbound_inputs, outputs


def write_saved_transform_from_session(
    session, inputs, outputs, export_path, as_text=False):
  """Write the current session as a SavedModel."""
  predict_signature_def = (
      tf.compat.v1.saved_model.signature_def_utils.predict_signature_def(
          inputs, outputs))

  builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_path)
  builder.add_meta_graph_and_variables(

      session, [constants.TRANSFORM_TAG],