# Path: tensorflow_transform/saved/constants.py
# TRANSFORM_TAG = 'transform'
# TRANSFORM_SIGNATURE = 'transform_signature'


from tensorflow_transform.saved import constants
from tensorflow.python.saved_model import loader_impl  # pylint: disable=g-direct-tensorflow-import
  Args:
    saved_model: A `SavedModel` protocol buffer.
    tags: Set of string tags to identify the required MetaGraphDef. These should
        correspond to the tags used when saving the variables using the
        SavedModel `save()` API.
  Returns:
    The chosen `MetaGraphDef` protocol buffer.  This can be used to further
    extract signature-defs, collection-defs, etc. If tags cannot be found,
    returns None.
  """
  result = None
  for meta_graph_def in saved_model.meta_graphs:
    if set(meta_graph_def.meta_info_def.tags) == set(tags):
      result = meta_graph_def
      break

  return result


def choose_meta_graph_def(saved_model):
  """Find a MetaGraphDef in the SavedModel with tag `constants.TRANSFORM_TAG`.

  Args:
    saved_model: A `SavedModel` protocol buffer.

  Returns:
    The chosen `MetaGraphDef` protocol buffer.  This can be used to further
    extract signature-defs, collection-defs, etc. If tags cannot be found,
    returns None.
  """

# Path: tensorflow_transform/saved/constants.py
# TRANSFORM_TAG = 'transform'
# TRANSFORM_SIGNATURE = 'transform_signature'


# Path: tensorflow_transform/saved/saved_model_loader.py
from tensorflow_transform.saved import constants
from tensorflow.python.saved_model import loader_impl  # pylint: disable=g-direct-tensorflow-import

  Args:
    saved_model: A `SavedModel` protocol buffer.
    tags: Set of string tags to identify the required MetaGraphDef. These should
        correspond to the tags used when saving the variables using the
        SavedModel `save()` API.
  Returns:
    The chosen `MetaGraphDef` protocol buffer.  This can be used to further
    extract signature-defs, collection-defs, etc. If tags cannot be found,
    returns None.
  """
  result = None
  for meta_graph_def in saved_model.meta_graphs:
    if set(meta_graph_def.meta_info_def.tags) == set(tags):
      result = meta_graph_def
      break

  return result


def choose_meta_graph_def(saved_model):
  """Find a MetaGraphDef in the SavedModel with tag `constants.TRANSFORM_TAG`.

  Args:
    saved_model: A `SavedModel` protocol buffer.

  Returns:
    The chosen `MetaGraphDef` protocol buffer.  This can be used to further
    extract signature-defs, collection-defs, etc. If tags cannot be found,
    returns None.
  """

  return _choose_meta_graph_def_internal(saved_model, [constants.TRANSFORM_TAG])