def _get_schema_annotations(graph, session):
  """Fetch extra_metadata annotations to be applied to the schema.

  Extracts any deferred annotations that have been added to the graph and
  evaluates them to obtain any_pb2.Any proto messages.

  Args:
    graph: A `tf.Graph` used to determine schema overrides.
    session: (optional) A `tf.Session` used to compute schema annotations.  If
      None, schema annotations will not be computed.

  Returns:
    tensor_annotations: dictionary from tensor to list of any_pb2.Any protos to
      be added as an annotation for that tensor's feature in the schema.
    global_annotations: list of any_pb2.Any protos to be added at the global
      schema level.
  """
  tensors = graph.get_collection(_TF_METADATA_EXTRA_ANNOTATION)
  type_urls = session.run(
      graph.get_collection(_TF_METADATA_EXTRA_ANNOTATION_TYPE_URL))
  proto_values = session.run(
      graph.get_collection(_TF_METADATA_EXTRA_ANNOTATION_PROTO))
  tensor_annotation_keys = []
  for tensor in tensors:
    # Entries meant for the global schema annotation will have names like
    # tft_schema_override_global_sentinel:0 or
    # transform/tft_schema_override_global_sentinel_1:0
    tensor_name = tensor.name.split('/')[-1]
    if tensor_name.startswith(_TF_METADATA_EXTRA_ANNOTATION_GLOBAL):
      tensor_annotation_keys.append(_TF_METADATA_EXTRA_ANNOTATION_GLOBAL)
    else:
      tensor_annotation_keys.append(tf_utils.hashable_tensor_or_op(tensor))
  return _get_schema_annotations_common(tensor_annotation_keys, type_urls,
                                        proto_values)
Fetch extra_metadata annotations to be applied to the schema.

  Extracts any deferred annotations that have been added to the graph and
  evaluates them to obtain any_pb2.Any proto messages.

  Args:
    graph: A `tf.Graph` used to determine schema overrides.
    session: (optional) A `tf.Session` used to compute schema annotations.  If
      None, schema annotations will not be computed.

  Returns:
    tensor_annotations: dictionary from tensor to list of any_pb2.Any protos to
      be added as an annotation for that tensor's feature in the schema.
    global_annotations: list of any_pb2.Any protos to be added at the global
      schema level.
  
tensorflow_transform/schema_inference.py