Merges the ASSETS_KEY collection into the GraphDefs in saved_model_proto.

  Removes the ASSETS_KEY collection from the GraphDefs in the SavedModel and
  modifies nodes with the assets filenames to point to the assets in `path`.
  After this transformation, the SavedModel GraphDefs can be used without
  feeding asset tensors.

  Args:
    saved_model_proto: SavedModel proto to be modified.
    path: path where the SavedModel is being loaded from.

def _merge_assets_key_collection(saved_model_proto, path):
  """Merges the ASSETS_KEY collection into the GraphDefs in saved_model_proto.

  Removes the ASSETS_KEY collection from the GraphDefs in the SavedModel and
  modifies nodes with the assets filenames to point to the assets in `path`.
  After this transformation, the SavedModel GraphDefs can be used without
  feeding asset tensors.

  Args:
    saved_model_proto: SavedModel proto to be modified.
    path: path where the SavedModel is being loaded from.
  """
  for meta_graph in saved_model_proto.meta_graphs:
    node_asset_map = {}
    if tf_v1.saved_model.constants.ASSETS_KEY in meta_graph.collection_def:
      assets_any_proto = meta_graph.collection_def[
          tf_v1.saved_model.constants.ASSETS_KEY].any_list.value
      for asset_any_proto in assets_any_proto:
        asset_proto = meta_graph_pb2.AssetFileDef()
        asset_any_proto.Unpack(asset_proto)
        asset_filename = _get_asset_filename(path, asset_proto.filename)
        node_asset_map[_get_node_name_from_tensor(
            asset_proto.tensor_info.name)] = asset_filename
      del meta_graph.collection_def[tf_v1.saved_model.constants.ASSETS_KEY]

    for node in meta_graph.graph_def.node:
      asset_filepath = node_asset_map.get(node.name)
      if asset_filepath:
        _check_asset_node_def(node)
        node.attr["value"].tensor.string_val[0] = asset_filepath