jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_export_graph_with_fixed_size_image_tensor_input(self):
    input_shape = [1, 320, 320, 3]

    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(
        trained_checkpoint_prefix, use_moving_averages=False)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel()
      output_directory = os.path.join(tmp_dir, 'output')
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory,
          input_shape=input_shape)
      saved_model_path = os.path.join(output_directory, 'saved_model')
      self.assertTrue(
          os.path.exists(os.path.join(saved_model_path, 'saved_model.pb')))

    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        meta_graph = tf.saved_model.loader.load(
            sess, [tf.saved_model.tag_constants.SERVING], saved_model_path)
        signature = meta_graph.signature_def['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        image_tensor = od_graph.get_tensor_by_name(input_tensor_name)
        self.assertSequenceEqual(image_tensor.get_shape().as_list(),
                                 input_shape)
//  [python]
def test_export_graph_with_fixed_size_image_tensor_input(self):
    input_shape = [1, 320, 320, 3]

    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(
        trained_checkpoint_prefix, use_moving_averages=False)
    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel()
      output_directory = os.path.join(tmp_dir, 'output')
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='image_tensor',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory,
          input_shape=input_shape)
      saved_model_path = os.path.join(output_directory, 'saved_model')
      self.assertTrue(
          os.path.exists(os.path.join(saved_model_path, 'saved_model.pb')))

    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        meta_graph = tf.saved_model.load(saved_model_path)
        signature = meta_graph.signatures['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        image_tensor = od_graph.get_tensor_by_name(input_tensor_name)
        self.assertSequenceEqual(image_tensor.get_shape().as_list(),
                                 input_shape)
//  [python]
def freeze_graph(input_saved_model_dir, output_saved_model_dir, node_names):
    with tf.Session() as sess:
        meta_graph_def = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], input_saved_model_dir)
        node_names = node_names + [n for n in meta_graph_def.collection_def["table_initializer"].node_list.value]
        node_names = node_names + [re.sub(r':\d$', "", n) for n in meta_graph_def.collection_def["asset_filepaths"].node_list.value]
        #frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, meta_graph_def.graph_def, node_names)
        frozen_graph_def = tf_graph_util.convert_variables_to_constants(sess, meta_graph_def.graph_def, node_names)

    if tf.gfile.IsDirectory(output_saved_model_dir):
        tf.gfile.DeleteRecursively(output_saved_model_dir)
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(frozen_graph_def, name="")
        if len(meta_graph_def.collection_def["table_initializer"].node_list.value) > 0:
            main_op = graph.get_operation_by_name(meta_graph_def.collection_def["table_initializer"].node_list.value[0])
        else:
            main_op = None
        assets = meta_graph_def.collection_def["asset_filepaths"].node_list.value
        assets = [graph.get_tensor_by_name(n) for n in assets]
        with tf.Session() as sess:
            builder = tf.saved_model.builder.SavedModelBuilder(output_saved_model_dir)
            builder.add_meta_graph_and_variables(
                    sess, [tf.saved_model.tag_constants.SERVING],
                    signature_def_map=meta_graph_def.signature_def,
                    assets_collection=assets,
                    main_op=main_op)
            builder.save()
//  [python]
def freeze_graph(input_saved_model_dir, output_saved_model_dir, node_names):
    with tf.Session() as sess:
        meta_graph_def = tf.saved_model.load(input_saved_model_dir, tags=[tf.saved_model.tag_constants.SERVING])
        node_names = node_names + [n for n in meta_graph_def.collection_def["table_initializer"].node_list.value]
        node_names = node_names + [re.sub(r':\d$', "", n) for n in meta_graph_def.collection_def["asset_filepaths"].node_list.value]
        #frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, meta_graph_def.graph_def, node_names)
        frozen_graph_def = tf_graph_util.convert_variables_to_constants(sess, meta_graph_def.graph_def, node_names)

    if tf.gfile.IsDirectory(output_saved_model_dir):
        tf.gfile.DeleteRecursively(output_saved_model_dir)
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(frozen_graph_def, name="")
        if len(meta_graph_def.collection_def["table_initializer"].node_list.value) > 0:
            main_op = graph.get_operation_by_name(meta_graph_def.collection_def["table_initializer"].node_list.value[0])
        else:
            main_op = None
        assets = meta_graph_def.collection_def["asset_filepaths"].node_list.value
        assets = [graph.get_tensor_by_name(n) for n in assets]
        with tf.Session() as sess:
            builder = tf.saved_model.builder.SavedModelBuilder(output_saved_model_dir)
            builder.add_meta_graph_and_variables(
                    sess, [tf.saved_model.tag_constants.SERVING],
                    signature_def_map=meta_graph_def.signature_def,
                    assets_collection=assets,
                    main_op=main_op)
            builder.save()
//  [python]
def export_saved_model(self, output_dir, rtol=None, atol=None, **kwargs):

        def _generate_signature():
            receiver_tensors = {
                'input_wids':
                tf.saved_model.utils.build_tensor_info(
                    self.pipeline.input_wids)
            }
            export_outputs = {
                'output_seqs':
                tf.saved_model.utils.build_tensor_info(
                    self.pipeline.output['output_seqs'])
            }

            signature_def = tf.saved_model.signature_def_utils.build_signature_def(
                receiver_tensors, export_outputs,
                tf.saved_model.signature_constants.PREDICT_METHOD_NAME)

            return {'translation_signature': signature_def}

        with self.pipeline._session.as_default() as sess:
            builder = tf.saved_model.builder.SavedModelBuilder(output_dir)
            builder.add_meta_graph_and_variables(
                sess, [tag_constants.SERVING],
                signature_def_map=_generate_signature(),
                assets_collection=ops.get_collection(
                    ops.GraphKeys.ASSET_FILEPATHS),
                clear_devices=True)
            builder.save()

        dummy_inputs = self.generate_dummy_inputs()
        with tf.Session(graph=tf.Graph()) as sess:
            # Restore model from the saved_model file, that is exported by TensorFlow estimator.
            MetaGraphDef = tf.saved_model.loader.load(sess, ['serve'],
                                                      output_dir)

            # SignatureDef protobuf
            SignatureDef_map = MetaGraphDef.signature_def
            SignatureDef = SignatureDef_map['translation_signature']
            # TensorInfo protobuf
            X_TensorInfo = SignatureDef.inputs['input_wids']
            y_TensorInfo = SignatureDef.outputs['output_seqs']
            X = tf.saved_model.utils.get_tensor_from_tensor_info(
                X_TensorInfo, sess.graph)
            y = tf.saved_model.utils.get_tensor_from_tensor_info(
                y_TensorInfo, sess.graph)
            outputs = sess.run(y, feed_dict={X: dummy_inputs['input_wids']})
            trans_result = self.pipeline.postprocess({'output_seqs': outputs})
            logger.info(trans_result)

        outputs_origin = self.pipeline.forward(
            {'input_ids': dummy_inputs['input_wids']})

        tols = {}
        if rtol is not None:
            tols['rtol'] = rtol
        if atol is not None:
            tols['atol'] = atol
        if not compare_arguments_nested('Output match failed', outputs,
                                        outputs_origin['output_seqs'], **tols):
            raise RuntimeError(
                'Export saved model failed because of validation error.')

        return {'model': output_dir}
//  [python]
def export_saved_model(self, output_dir, rtol=None, atol=None, **kwargs):

        def _generate_signature():
            receiver_tensors = {
                'input_wids':
                tf.saved_model.utils.build_tensor_info(
                    self.pipeline.input_wids)
            }
            export_outputs = {
                'output_seqs':
                tf.saved_model.utils.build_tensor_info(
                    self.pipeline.output['output_seqs'])
            }

            signature_def = tf.saved_model.signature_def_utils.build_signature_def(
                receiver_tensors, export_outputs,
                tf.saved_model.signature_constants.PREDICT_METHOD_NAME)

            return {'translation_signature': signature_def}

        with self.pipeline._session.as_default() as sess:
            builder = tf.saved_model.builder.SavedModelBuilder(output_dir)
            builder.add_meta_graph_and_variables(
                sess, [tag_constants.SERVING],
                signature_def_map=_generate_signature(),
                assets_collection=ops.get_collection(
                    ops.GraphKeys.ASSET_FILEPATHS),
                clear_devices=True)
            builder.save()

        dummy_inputs = self.generate_dummy_inputs()
        with tf.Session(graph=tf.Graph()) as sess:
            # Restore model from the saved_model file, that is exported by TensorFlow estimator.
            MetaGraphDef = tf.saved_model.load(sess, output_dir)

            # SignatureDef protobuf
            SignatureDef_map = MetaGraphDef.signature_def
            SignatureDef = SignatureDef_map['translation_signature']
            # TensorInfo protobuf
            X_TensorInfo = SignatureDef.inputs['input_wids']
            y_TensorInfo = SignatureDef.outputs['output_seqs']
            X = tf.saved_model.utils.get_tensor_from_tensor_info(
                X_TensorInfo, sess.graph)
            y = tf.saved_model.utils.get_tensor_from_tensor_info(
                y_TensorInfo, sess.graph)
            outputs = sess.run(y, feed_dict={X: dummy_inputs['input_wids']})
            trans_result = self.pipeline.postprocess({'output_seqs': outputs})
            logger.info(trans_result)

        outputs_origin = self.pipeline.forward(
            {'input_ids': dummy_inputs['input_wids']})

        tols = {}
        if rtol is not None:
            tols['rtol'] = rtol
        if atol is not None:
            tols['atol'] = atol
        if not compare_arguments_nested('Output match failed', outputs,
                                        outputs_origin['output_seqs'], **tols):
            raise RuntimeError(
                'Export saved model failed because of validation error.')

        return {'model': output_dir}
//  [python]
def _from_saved_model_v1(sess, model_path, input_names, output_names, tag, signature_names, use_graph_names):
    """Load tensorflow graph from saved_model."""

    wrn_no_tag = "'--tag' not specified for saved_model. Using --tag serve"
    wrn_empty_tag = "'--tag' value is empty string. Using tags = []"
    wrn_empty_sig = "'--signature_def' not provided. Using all signatures."

    if tag is None:
        tag = [tf.saved_model.tag_constants.SERVING]
        logger.warning(wrn_no_tag)

    if not signature_names:
        logger.warning(wrn_empty_sig)

    if tag == '':
        tag = []
        logger.warning(wrn_empty_tag)

    if not isinstance(tag, list):
        tag = [tag]

    imported = tf.saved_model.loader.load(sess, tag, model_path)
    signatures = []
    for k in imported.signature_def.keys():
        if k in signature_names or (not signature_names and not k.startswith("_")):
            signatures.append(k)
    try:
        from tensorflow.contrib.saved_model.python.saved_model import signature_def_utils
        # pylint: disable=unnecessary-lambda
        get_signature_def = lambda meta_graph_def, k: \
            signature_def_utils.get_signature_def_by_key(meta_graph_def, k)
    except ImportError:
        # TF1.12 changed the api
        get_signature_def = lambda meta_graph_def, k: meta_graph_def.signature_def[k]

    tensors_to_rename = {}
    if input_names is None:
        input_names = []
        for k in signatures:
            inputs_tensor_info = get_signature_def(imported, k).inputs
            for structured_name, input_tensor in inputs_tensor_info.items():
                if input_tensor.name not in input_names:
                    input_names.append(input_tensor.name)
                    if not use_graph_names:
                        tensors_to_rename[input_tensor.name] = structured_name
    if output_names is None:
        output_names = []
        for k in signatures:
            outputs_tensor_info = get_signature_def(imported, k).outputs
            for structured_name, output_tensor in outputs_tensor_info.items():
                if output_tensor.name not in output_names:
                    output_names.append(output_tensor.name)
                    if not use_graph_names:
                        tensors_to_rename[output_tensor.name] = structured_name
    frozen_graph, initialized_tables = \
        freeze_session(sess, input_names=input_names, output_names=output_names, get_tables=True)
    return frozen_graph, input_names, output_names, initialized_tables, tensors_to_rename
//  [python]
def _from_saved_model_v1(sess, model_path, input_names, output_names, tag, signature_names, use_graph_names):
    """Load tensorflow graph from saved_model."""

    wrn_no_tag = "'--tag' not specified for saved_model. Using --tag serve"
    wrn_empty_tag = "'--tag' value is empty string. Using tags = []"
    wrn_empty_sig = "'--signature_def' not provided. Using all signatures."

    if tag is None:
        tag = [tf.saved_model.tag_constants.SERVING]
        logger.warning(wrn_no_tag)

    if not signature_names:
        logger.warning(wrn_empty_sig)

    if tag == '':
        tag = []
        logger.warning(wrn_empty_tag)

    if not isinstance(tag, list):
        tag = [tag]

    imported = tf.saved_model.load(sess, model_path, tags=tag)
    signatures = []
    for k in imported.signature_def.keys():
        if k in signature_names or (not signature_names and not k.startswith("_")):
            signatures.append(k)
    try:
        from tensorflow.contrib.saved_model.python.saved_model import signature_def_utils
        # pylint: disable=unnecessary-lambda
        get_signature_def = lambda meta_graph_def, k: \
            signature_def_utils.get_signature_def_by_key(meta_graph_def, k)
    except ImportError:
        # TF1.12 changed the api
        get_signature_def = lambda meta_graph_def, k: meta_graph_def.signature_def[k]

    tensors_to_rename = {}
    if input_names is None:
        input_names = []
        for k in signatures:
            inputs_tensor_info = get_signature_def(imported, k).inputs
            for structured_name, input_tensor in inputs_tensor_info.items():
                if input_tensor.name not in input_names:
                    input_names.append(input_tensor.name)
                    if not use_graph_names:
                        tensors_to_rename[input_tensor.name] = structured_name
    if output_names is None:
        output_names = []
        for k in signatures:
            outputs_tensor_info = get_signature_def(imported, k).outputs
            for structured_name, output_tensor in outputs_tensor_info.items():
                if output_tensor.name not in output_names:
                    output_names.append(output_tensor.name)
                    if not use_graph_names:
                        tensors_to_rename[output_tensor.name] = structured_name
    frozen_graph, initialized_tables = \
        freeze_session(sess, input_names=input_names, output_names=output_names, get_tables=True)
    return frozen_graph, input_names, output_names, initialized_tables, tensors_to_rename
//  [python]
def test_export_saved_model_and_run_inference(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
    output_directory = os.path.join(tmp_dir, 'output')
    saved_model_path = os.path.join(output_directory, 'saved_model')

    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='tf_example',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)

    tf_example_np = np.hstack([self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8))] * 2)
    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        meta_graph = tf.saved_model.loader.load(
            sess, [tf.saved_model.tag_constants.SERVING], saved_model_path)

        signature = meta_graph.signature_def['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        tf_example = od_graph.get_tensor_by_name(input_tensor_name)

        boxes = od_graph.get_tensor_by_name(
            signature.outputs['detection_boxes'].name)
        scores = od_graph.get_tensor_by_name(
            signature.outputs['detection_scores'].name)
        classes = od_graph.get_tensor_by_name(
            signature.outputs['detection_classes'].name)
        keypoints = od_graph.get_tensor_by_name(
            signature.outputs['detection_keypoints'].name)
        masks = od_graph.get_tensor_by_name(
            signature.outputs['detection_masks'].name)
        num_detections = od_graph.get_tensor_by_name(
            signature.outputs['num_detections'].name)

        (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
         num_detections_np) = sess.run(
             [boxes, scores, classes, keypoints, masks, num_detections],
             feed_dict={tf_example: tf_example_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])
//  [python]
def test_export_saved_model_and_run_inference(self):
    tmp_dir = self.get_temp_dir()
    trained_checkpoint_prefix = os.path.join(tmp_dir, 'model.ckpt')
    self._save_checkpoint_from_mock_model(trained_checkpoint_prefix,
                                          use_moving_averages=False)
    output_directory = os.path.join(tmp_dir, 'output')
    saved_model_path = os.path.join(output_directory, 'saved_model')

    with mock.patch.object(
        model_builder, 'build', autospec=True) as mock_builder:
      mock_builder.return_value = FakeModel(
          add_detection_keypoints=True, add_detection_masks=True)
      pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
      pipeline_config.eval_config.use_moving_averages = False
      exporter.export_inference_graph(
          input_type='tf_example',
          pipeline_config=pipeline_config,
          trained_checkpoint_prefix=trained_checkpoint_prefix,
          output_directory=output_directory)

    tf_example_np = np.hstack([self._create_tf_example(
        np.ones((4, 4, 3)).astype(np.uint8))] * 2)
    with tf.Graph().as_default() as od_graph:
      with self.test_session(graph=od_graph) as sess:
        meta_graph = tf.saved_model.load(
            saved_model_path, tags=[tf.saved_model.tag_constants.SERVING])

        signature = meta_graph.signature_def['serving_default']
        input_tensor_name = signature.inputs['inputs'].name
        tf_example = od_graph.get_tensor_by_name(input_tensor_name)

        boxes = od_graph.get_tensor_by_name(
            signature.outputs['detection_boxes'].name)
        scores = od_graph.get_tensor_by_name(
            signature.outputs['detection_scores'].name)
        classes = od_graph.get_tensor_by_name(
            signature.outputs['detection_classes'].name)
        keypoints = od_graph.get_tensor_by_name(
            signature.outputs['detection_keypoints'].name)
        masks = od_graph.get_tensor_by_name(
            signature.outputs['detection_masks'].name)
        num_detections = od_graph.get_tensor_by_name(
            signature.outputs['num_detections'].name)

        (boxes_np, scores_np, classes_np, keypoints_np, masks_np,
         num_detections_np) = sess.run(
             [boxes, scores, classes, keypoints, masks, num_detections],
             feed_dict={tf_example: tf_example_np})
        self.assertAllClose(boxes_np, [[[0.0, 0.0, 0.5, 0.5],
                                        [0.5, 0.5, 0.8, 0.8]],
                                       [[0.5, 0.5, 1.0, 1.0],
                                        [0.0, 0.0, 0.0, 0.0]]])
        self.assertAllClose(scores_np, [[0.7, 0.6],
                                        [0.9, 0.0]])
        self.assertAllClose(classes_np, [[1, 2],
                                         [2, 1]])
        self.assertAllClose(keypoints_np, np.arange(48).reshape([2, 2, 6, 2]))
        self.assertAllClose(masks_np, np.arange(64).reshape([2, 2, 4, 4]))
        self.assertAllClose(num_detections_np, [2, 1])
//  [python]
def load(self, saved_model_dir_or_frozen_graph: Text):
    """Load the model using saved model or a frozen graph."""
    if not self.sess:
      self.sess = self._build_session()
    self.signitures = {
        'image_files': 'image_files:0',
        'image_arrays': 'image_arrays:0',
        'prediction': 'detections:0',
    }

    # Load saved model if it is a folder.
    if tf.io.gfile.isdir(saved_model_dir_or_frozen_graph):
      return tf.saved_model.load(self.sess, ['serve'],
                                 saved_model_dir_or_frozen_graph)

    # Load a frozen graph.
    graph_def = tf.GraphDef()
    with tf.gfile.GFile(saved_model_dir_or_frozen_graph, 'rb') as f:
      graph_def.ParseFromString(f.read())
    return tf.import_graph_def(graph_def, name='')
//  [python]
def load(self, saved_model_dir_or_frozen_graph: Text):
    """Load the model using saved model or a frozen graph."""
    if not self.sess:
      self.sess = self._build_session()
    self.signitures = {
        'image_files': 'image_files:0',
        'image_arrays': 'image_arrays:0',
        'prediction': 'detections:0',
    }

    # Load saved model if it is a folder.
    if tf.io.gfile.isdir(saved_model_dir_or_frozen_graph):
      return tf.saved_model.load(self.sess, tags=['serve'],
                                 export_dir=saved_model_dir_or_frozen_graph)

    # Load a frozen graph.
    graph_def = tf.GraphDef()
    with tf.gfile.GFile(saved_model_dir_or_frozen_graph, 'rb') as f:
      graph_def.ParseFromString(f.read())
    return tf.import_graph_def(graph_def, name='')
//  [python]
def MakeExtractor(sess, config, import_scope=None):
    """Creates a function to extract features from an image.

  Args:
    sess: TensorFlow session to use.
    config: DelfConfig proto containing the model configuration.
    import_scope: Optional scope to use for model.

  Returns:
    Function that receives an image and returns features.
  """
    tf.saved_model.loader.load(
        sess, [tf.saved_model.tag_constants.SERVING],
        config.model_path,
        import_scope=import_scope)
    import_scope_prefix = import_scope + '/' if import_scope is not None else ''
    input_image = sess.graph.get_tensor_by_name('%sinput_image:0' %
                                                import_scope_prefix)
    input_score_threshold = sess.graph.get_tensor_by_name(
        '%sinput_abs_thres:0' % import_scope_prefix)
    input_image_scales = sess.graph.get_tensor_by_name('%sinput_scales:0' %
                                                       import_scope_prefix)
    input_max_feature_num = sess.graph.get_tensor_by_name(
        '%sinput_max_feature_num:0' % import_scope_prefix)
    boxes = sess.graph.get_tensor_by_name('%sboxes:0' % import_scope_prefix)
    raw_descriptors = sess.graph.get_tensor_by_name('%sfeatures:0' %
                                                    import_scope_prefix)
    feature_scales = sess.graph.get_tensor_by_name('%sscales:0' %
                                                   import_scope_prefix)
    attention_with_extra_dim = sess.graph.get_tensor_by_name(
        '%sscores:0' % import_scope_prefix)
    attention = tf.reshape(attention_with_extra_dim,
                           [tf.shape(attention_with_extra_dim)[0]])

    locations, descriptors = feature_extractor.DelfFeaturePostProcessing(
        boxes, raw_descriptors, config)

    def ExtractorFn(image):
        """Receives an image and returns DELF features.

    Args:
      image: Uint8 array with shape (height, width 3) containing the RGB image.

    Returns:
      Tuple (locations, descriptors, feature_scales, attention)
    """
        return sess.run([locations, descriptors, feature_scales, attention],
                        feed_dict={
                            input_image: image,
                            input_score_threshold: config.delf_local_config.score_threshold,
                            input_image_scales: list(config.image_scales),
                            input_max_feature_num: config.delf_local_config.max_feature_num
                        })

    return ExtractorFn
//  [python]
def MakeExtractor(sess, config, import_scope=None):
    """Creates a function to extract features from an image.

  Args:
    sess: TensorFlow session to use.
    config: DelfConfig proto containing the model configuration.
    import_scope: Optional scope to use for model.

  Returns:
    Function that receives an image and returns features.
  """
    tf.saved_model.load(
        sess, [tf.saved_model.tag_constants.SERVING],
        config.model_path,
        import_scope=import_scope)
    import_scope_prefix = import_scope + '/' if import_scope is not None else ''
    input_image = sess.graph.get_tensor_by_name('%sinput_image:0' %
                                                import_scope_prefix)
    input_score_threshold = sess.graph.get_tensor_by_name(
        '%sinput_abs_thres:0' % import_scope_prefix)
    input_image_scales = sess.graph.get_tensor_by_name('%sinput_scales:0' %
                                                       import_scope_prefix)
    input_max_feature_num = sess.graph.get_tensor_by_name(
        '%sinput_max_feature_num:0' % import_scope_prefix)
    boxes = sess.graph.get_tensor_by_name('%sboxes:0' % import_scope_prefix)
    raw_descriptors = sess.graph.get_tensor_by_name('%sfeatures:0' %
                                                    import_scope_prefix)
    feature_scales = sess.graph.get_tensor_by_name('%sscales:0' %
                                                   import_scope_prefix)
    attention_with_extra_dim = sess.graph.get_tensor_by_name(
        '%sscores:0' % import_scope_prefix)
    attention = tf.reshape(attention_with_extra_dim,
                           [tf.shape(attention_with_extra_dim)[0]])

    locations, descriptors = feature_extractor.DelfFeaturePostProcessing(
        boxes, raw_descriptors, config)

    def ExtractorFn(image):
        """Receives an image and returns DELF features.

    Args:
      image: Uint8 array with shape (height, width 3) containing the RGB image.

    Returns:
      Tuple (locations, descriptors, feature_scales, attention)
    """
        return sess.run([locations, descriptors, feature_scales, attention],
                        feed_dict={
                            input_image: image,
                            input_score_threshold: config.delf_local_config.score_threshold,
                            input_image_scales: list(config.image_scales),
                            input_max_feature_num: config.delf_local_config.max_feature_num
                        })

    return ExtractorFn
//  [python]
def main(unused_argv):

  assert FLAGS.output_dir, 'Must specify --output_dir.'
  output_dir = FLAGS.output_dir
  if not tf.gfile.Exists(output_dir):
    tf.gfile.MakeDirs(output_dir)

  with tf.Session(graph=tf.Graph()) as sess:
    print(' - Loading saved model...')
    tf.saved_model.loader.load(sess, FLAGS.tag_set.split(','),
                               FLAGS.saved_model_dir)

    image_files = tf.gfile.Glob(FLAGS.image_file_pattern)
    for i, image_file in enumerate(image_files):
      print(' - Processing image %d...' % i)

      if FLAGS.compressed_input:
        options = tf.python_io.TFRecordOptions(
            tf.python_io.TFRecordCompressionType.GZIP)
      else:
        options = tf.python_io.TFRecordOptions()
      record_iterator = tf.python_io.tf_record_iterator(
          path=image_file, options=options)
      np_input = None
      if FLAGS.input_type == 'tf_example':
        for tf_example_string in record_iterator:
          image_string = tf_example_string
          break  # Only read the first tf.Example.
        np_input = prepare_input_feed(image_string, FLAGS.input_type)
      elif FLAGS.input_type == 'image_tensor':
        example = tf.train.Example()
        for tf_example_string in record_iterator:
          example.ParseFromString(tf_example_string)
          break  # Only read the first tf.Example.
        f = example.features.feature
        image = f['image/ct_image'].bytes_list.value[0]
        image = np.fromstring(
            image, dtype=np.float32).reshape(DESIRED_IMAGE_SHAPE)
        np_input = prepare_input_feed(image, FLAGS.input_type)
      else:
        raise ValueError('Unknown input type %s' % FLAGS.input_type)

      output_nodes = [
          FLAGS.output_classes_node,
          FLAGS.output_scores_node,
      ]
      outputs = sess.run(output_nodes, feed_dict={FLAGS.input_node: np_input})
      classes = outputs[0]
      scores = outputs[1]
      scores = np.squeeze(scores, axis=(0,))
      classes = np.squeeze(classes.astype(np.int32), axis=(0,))

      output_path = os.path.join(output_dir,
                                 os.path.basename(image_file) + '.npz')
      with tf.gfile.GFile(output_path, 'w') as f:
        io_buffer = io.BytesIO()
        np.savez(io_buffer, scores=scores, classes=classes)
        f.write(io_buffer.getvalue())
      print('Output to %s' % output_path)

  print(' - Done!')
//  [python]
def main(unused_argv):

  assert FLAGS.output_dir, 'Must specify --output_dir.'
  output_dir = FLAGS.output_dir
  if not tf.gfile.Exists(output_dir):
    tf.gfile.MakeDirs(output_dir)

  with tf.Session(graph=tf.Graph()) as sess:
    print(' - Loading saved model...')
    tf.saved_model.load(sess, FLAGS.saved_model_dir, tags=FLAGS.tag_set.split(','))

    image_files = tf.gfile.Glob(FLAGS.image_file_pattern)
    for i, image_file in enumerate(image_files):
      print(' - Processing image %d...' % i)

      if FLAGS.compressed_input:
        options = tf.python_io.TFRecordOptions(
            tf.python_io.TFRecordCompressionType.GZIP)
      else:
        options = tf.python_io.TFRecordOptions()
      record_iterator = tf.python_io.tf_record_iterator(
          path=image_file, options=options)
      np_input = None
      if FLAGS.input_type == 'tf_example':
        for tf_example_string in record_iterator:
          image_string = tf_example_string
          break  # Only read the first tf.Example.
        np_input = prepare_input_feed(image_string, FLAGS.input_type)
      elif FLAGS.input_type == 'image_tensor':
        example = tf.train.Example()
        for tf_example_string in record_iterator:
          example.ParseFromString(tf_example_string)
          break  # Only read the first tf.Example.
        f = example.features.feature
        image = f['image/ct_image'].bytes_list.value[0]
        image = np.fromstring(
            image, dtype=np.float32).reshape(DESIRED_IMAGE_SHAPE)
        np_input = prepare_input_feed(image, FLAGS.input_type)
      else:
        raise ValueError('Unknown input type %s' % FLAGS.input_type)

      output_nodes = [
          FLAGS.output_classes_node,
          FLAGS.output_scores_node,
      ]
      outputs = sess.run(output_nodes, feed_dict={FLAGS.input_node: np_input})
      classes = outputs[0]
      scores = outputs[1]
      scores = np.squeeze(scores, axis=(0,))
      classes = np.squeeze(classes.astype(np.int32), axis=(0,))

      output_path = os.path.join(output_dir,
                                 os.path.basename(image_file) + '.npz')
      with tf.gfile.GFile(output_path, 'w') as f:
        io_buffer = io.BytesIO()
        np.savez(io_buffer, scores=scores, classes=classes)
        f.write(io_buffer.getvalue())
      print('Output to %s' % output_path)

  print(' - Done!')
//  [python]
def _load_tensorflow_saved_model(tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key,
                                 tf_sess=None):
    """
    Load a specified TensorFlow model consisting of a TensorFlow metagraph and signature definition
    from a serialized TensorFlow ``SavedModel`` collection.

    :param tf_saved_model_dir: The local filesystem path or run-relative artifact path to the model.
    :param tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
                               serialized ``SavedModel`` object. For more information, see the
                               ``tags`` parameter of the `tf.saved_model.builder.SavedModelBuilder
                               method <https://www.tensorflow.org/api_docs/python/tf/saved_model/
                               builder/SavedModelBuilder#add_meta_graph>`_.
    :param tf_signature_def_key: A string identifying the input/output signature associated with the
                                 model. This is a key within the serialized ``SavedModel``'s
                                 signature definition mapping. For more information, see the
                                 ``signature_def_map`` parameter of the
                                 ``tf.saved_model.builder.SavedModelBuilder`` method.
    :param tf_sess: The TensorFlow session in which to load the metagraph.
                    Required in TensorFlow versions < 2.0.0. Unused in TensorFlow versions >= 2.0.0
    :return: For TensorFlow versions < 2.0.0:
             A TensorFlow signature definition of type:
             ``tensorflow.core.protobuf.meta_graph_pb2.SignatureDef``. This defines input and
             output tensors within the specified metagraph for inference.
             For TensorFlow versions >= 2.0.0:
             A callable graph (tensorflow.function) that takes inputs and returns inferences.
    """
    if LooseVersion(tensorflow.__version__) < LooseVersion('2.0.0'):
        loaded = tensorflow.saved_model.loader.load(
            sess=tf_sess,
            tags=tf_meta_graph_tags,
            export_dir=tf_saved_model_dir)
        loaded_sig = loaded.signature_def
    else:
        loaded = tensorflow.saved_model.load(  # pylint: disable=no-value-for-parameter
                tags=tf_meta_graph_tags,
                export_dir=tf_saved_model_dir)
        loaded_sig = loaded.signatures
    if tf_signature_def_key not in loaded_sig:
        raise MlflowException("Could not find signature def key %s. Available keys are: %s"
                              % (tf_signature_def_key, list(loaded_sig.keys())))
    return loaded_sig[tf_signature_def_key]
//  [python]
def _load_tensorflow_saved_model(tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key,
                                 tf_sess=None):
    """
    Load a specified TensorFlow model consisting of a TensorFlow metagraph and signature definition
    from a serialized TensorFlow ``SavedModel`` collection.

    :param tf_saved_model_dir: The local filesystem path or run-relative artifact path to the model.
    :param tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
                               serialized ``SavedModel`` object. For more information, see the
                               ``tags`` parameter of the `tf.saved_model.builder.SavedModelBuilder
                               method <https://www.tensorflow.org/api_docs/python/tf/saved_model/
                               builder/SavedModelBuilder#add_meta_graph>`_.
    :param tf_signature_def_key: A string identifying the input/output signature associated with the
                                 model. This is a key within the serialized ``SavedModel``'s
                                 signature definition mapping. For more information, see the
                                 ``signature_def_map`` parameter of the
                                 ``tf.saved_model.builder.SavedModelBuilder`` method.
    :param tf_sess: The TensorFlow session in which to load the metagraph.
                    Required in TensorFlow versions < 2.0.0. Unused in TensorFlow versions >= 2.0.0
    :return: For TensorFlow versions < 2.0.0:
             A TensorFlow signature definition of type:
             ``tensorflow.core.protobuf.meta_graph_pb2.SignatureDef``. This defines input and
             output tensors within the specified metagraph for inference.
             For TensorFlow versions >= 2.0.0:
             A callable graph (tensorflow.function) that takes inputs and returns inferences.
    """
    loaded = tensorflow.saved_model.load(  # pylint: disable=no-value-for-parameter
                tags=tf_meta_graph_tags,
                export_dir=tf_saved_model_dir)
    loaded_sig = loaded.signatures
    if tf_signature_def_key not in loaded_sig:
        raise MlflowException("Could not find signature def key %s. Available keys are: %s"
                              % (tf_signature_def_key, list(loaded_sig.keys())))
    return loaded_sig[tf_signature_def_key]
//  [python]
def calculate_style_bottleneck(session,
                               predict_saved_model,
                               style_dataset,
                               batch_size=10,
                               min_sample_count=30):
  """Calculate style bottleneck using style predict SavedModel.

  Args:
    session: tf.Session, the session that contains subgraph to load the training
      dataset.
    predict_saved_model: str, path to the style predict SavedModel.
    style_dataset: tf.data.Dataset that contains training style images.
    batch_size: int, Batch size used when calculating style bottlenecks.
    min_sample_count: int, the minimum number of sample needed.

  Returns:
    [[[int]]] Value of style bottleneck for the style images being processed.
  """

  styles_bottleneck = None

  tf.logging.info('Started calculating style bottlenecks.')

  # Load and run the style predict SavedModel
  loaded = tf.saved_model.loader.load(session, ['serve'], predict_saved_model)
  bottleneck_tensor_name = loaded.signature_def['serving_default'].outputs[
      'style_bottleneck'].name

  # Load the style dataset
  dataset = style_dataset.batch(batch_size)
  iterator = tf.data.make_initializable_iterator(dataset)
  session.run(iterator.initializer)
  next_element = iterator.get_next()

  for _ in range(np.math.ceil(min_sample_count / batch_size)):
    style_image = session.run(next_element)

    style_bottleneck_batch = session.run(
        bottleneck_tensor_name, feed_dict={'style_image:0': style_image})
    if styles_bottleneck is None:
      styles_bottleneck = style_bottleneck_batch
    else:
      styles_bottleneck = np.append(
          styles_bottleneck, style_bottleneck_batch, axis=0)

  tf.logging.info('Finished calculating style bottlenecks.')

  return styles_bottleneck
//  [python]
def calculate_style_bottleneck(session,
                               predict_saved_model,
                               style_dataset,
                               batch_size=10,
                               min_sample_count=30):
  """Calculate style bottleneck using style predict SavedModel.

  Args:
    session: tf.Session, the session that contains subgraph to load the training
      dataset.
    predict_saved_model: str, path to the style predict SavedModel.
    style_dataset: tf.data.Dataset that contains training style images.
    batch_size: int, Batch size used when calculating style bottlenecks.
    min_sample_count: int, the minimum number of sample needed.

  Returns:
    [[[int]]] Value of style bottleneck for the style images being processed.
  """

  styles_bottleneck = None

  tf.logging.info('Started calculating style bottlenecks.')

  # Load and run the style predict SavedModel
  loaded = tf.saved_model.load(predict_saved_model)
  bottleneck_tensor_name = loaded.signatures['serving_default'].outputs[
      'style_bottleneck'].name

  # Load the style dataset
  dataset = style_dataset.batch(batch_size)
  iterator = tf.data.make_initializable_iterator(dataset)
  session.run(iterator.initializer)
  next_element = iterator.get_next()

  for _ in range(np.math.ceil(min_sample_count / batch_size)):
    style_image = session.run(next_element)

    style_bottleneck_batch = session.run(
        bottleneck_tensor_name, feed_dict={'style_image:0': style_image})
    if styles_bottleneck is None:
      styles_bottleneck = style_bottleneck_batch
    else:
      styles_bottleneck = np.append(
          styles_bottleneck, style_bottleneck_batch, axis=0)

  tf.logging.info('Finished calculating style bottlenecks.')

  return styles_bottleneck
//  [python]
def getSess(self):
        flags = tf.flags
        FLAGS = flags.FLAGS

        flags.DEFINE_string(
            "export_dir", '/Users/shiyu/Documents/Project/intent_bot_working/bert_small_10epoch/',
            "The dir where the exported model has been written.")

        tokenizer = tokenization.FullTokenizer(vocab_file='/Users/shiyu/Documents/Project/intent_bot_working/bert_small_10epoch/vocab.txt',
                                               do_lower_case=True)
        sess = tf.InteractiveSession()
        tf.saved_model.loader.load(sess, [tag_constants.SERVING], FLAGS.export_dir)
        return sess, tokenizer
//  [python]
def getSess(self):
        flags = tf.flags
        FLAGS = flags.FLAGS

        flags.DEFINE_string(
            "export_dir", '/Users/shiyu/Documents/Project/intent_bot_working/bert_small_10epoch/',
            "The dir where the exported model has been written.")

        tokenizer = tokenization.FullTokenizer(vocab_file='/Users/shiyu/Documents/Project/intent_bot_working/bert_small_10epoch/vocab.txt',
                                               do_lower_case=True)
        sess = tf.compat.v1.InteractiveSession()
        tf.saved_model.load(sess, FLAGS.export_dir)
        return sess, tokenizer
//  [python]
def _make_rpc_friendly(model_dir, tag, in_map, out_map, vars_path):
  """Wraps a SavedModel in EE RPC-friendly ops and saves a temporary copy."""
  out_dir = tempfile.mkdtemp()
  builder = tf.saved_model.Builder(out_dir)

  # Get a GraphDef from the saved model
  with tf.Session() as sesh:
    meta_graph = tf.saved_model.load(sesh, [tag], model_dir)

  graph_def = meta_graph.graph_def

  # Purge the default graph immediately after: we want to remap parts of the
  # graph when we load it and we don't know what those parts are yet.
  tf.reset_default_graph()

  input_op_keys = list(in_map.keys())
  input_new_keys = list(in_map.values())

  # Get the shape and type of the input tensors
  in_op_types = _get_input_tensor_spec(graph_def, input_op_keys)

  # Create new input placeholders to receive RPC TensorProto payloads
  in_op_map = {
      k: tf.placeholder(
          tf.string, shape=[None], name='earthengine_in_{}'.format(i))
      for (i, k) in enumerate(input_new_keys)
  }

  # Glue on decoding ops to remap to the imported graph.
  decoded_op_map = {
      k: _decode_op(in_op_map[in_map[k]], in_op_types[k])
      for k in input_op_keys
  }

  # Okay now we're ready to import the graph again but remapped.
  saver = tf.train.import_meta_graph(
      meta_graph_or_file=meta_graph, input_map=decoded_op_map)

  # Boilerplate to build a signature def for our new graph
  sig_in = {
      _strip_index(k):
      saved_model_utils.build_tensor_info(v) for (k, v) in in_op_map.items()
  }

  sig_out = {}
  for index, (k, v) in enumerate(out_map.items()):
    out_tensor = saved_model_utils.build_tensor_info(
        _encode_op(
            tf.get_default_graph().get_tensor_by_name(k),
            name='earthengine_out_{}'.format(index)))

    sig_out[_strip_index(v)] = out_tensor

  sig_def = signature_def_utils.build_signature_def(
      sig_in, sig_out, signature_constants.PREDICT_METHOD_NAME)

  # Open a new session to load the variables and add them to the builder.
  with tf.Session() as sesh:
    if saver:
      saver.restore(sesh, model_dir + vars_path)
    builder.add_meta_graph_and_variables(
        sesh,
        tags=[tf.saved_model.tag_constants.SERVING],
        signature_def_map={
            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: sig_def
        },
        saver=saver)

  builder.save()
  return out_dir