def test_build_all_signature_defs(self):
    input_features = constant_op.constant(["10"])
    input_example = constant_op.constant(["11"])
    input_ops = input_fn_utils.InputFnOps({
        "features": input_features
    }, None, {"default input": input_example})
    input_alternatives, _ = (
        saved_model_export_utils.get_input_alternatives(input_ops))
    output_1 = constant_op.constant(["1"])
    output_2 = constant_op.constant(["2"])
    output_3 = constant_op.constant(["3"])
    provided_output_alternatives = {
        "head-1": (constants.ProblemType.LINEAR_REGRESSION, {
            "some_output_1": output_1
        }),
        "head-2": (constants.ProblemType.CLASSIFICATION, {
            "some_output_2": output_2
        }),
        "head-3": (constants.ProblemType.UNSPECIFIED, {
            "some_output_3": output_3
        }),
    }
    model_fn_ops = model_fn.ModelFnOps(
        model_fn.ModeKeys.INFER,
        predictions={"some_output": constant_op.constant(["4"])},
        output_alternatives=provided_output_alternatives)
    output_alternatives, _ = (saved_model_export_utils.get_output_alternatives(
        model_fn_ops, "head-1"))

    signature_defs = saved_model_export_utils.build_all_signature_defs(
        input_alternatives, output_alternatives, "head-1")

    expected_signature_defs = {
        "serving_default":
            signature_def_utils.regression_signature_def(input_example,
                                                         output_1),
        "default_input_alternative:head-1":
            signature_def_utils.regression_signature_def(input_example,
                                                         output_1),
        "default_input_alternative:head-2":
            signature_def_utils.classification_signature_def(input_example,
                                                             output_2, None),
        "default_input_alternative:head-3":
            signature_def_utils.predict_signature_def({
                "input": input_example
            }, {"output": output_3}),
        "features_input_alternative:head-1":
            signature_def_utils.regression_signature_def(input_features,
                                                         output_1),
        "features_input_alternative:head-2":
            signature_def_utils.classification_signature_def(input_features,
                                                             output_2, None),
        "features_input_alternative:head-3":
            signature_def_utils.predict_signature_def({
                "input": input_features
            }, {"output": output_3}),
    }

    self.assertDictEqual(expected_signature_defs, signature_defs)

  def test_build_all_signature_defs_without_receiver_alternatives(self):
    receiver_tensor = array_ops.placeholder(dtypes.string)
    output_1 = constant_op.constant([1.])
    output_2 = constant_op.constant(["2"])
    output_3 = constant_op.constant(["3"])
    export_outputs = {
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
            export_output.RegressionOutput(value=output_1),
        "head-2": export_output.ClassificationOutput(classes=output_2),
        "head-3": export_output.PredictOutput(outputs={
            "some_output_3": output_3
        }),
    }

    signature_defs = export.build_all_signature_defs(
        receiver_tensor, export_outputs)

    expected_signature_defs = {
        "serving_default":
            signature_def_utils.regression_signature_def(receiver_tensor,
                                                         output_1),
        "head-2":
            signature_def_utils.classification_signature_def(receiver_tensor,
                                                             output_2, None),
        "head-3":
            signature_def_utils.predict_signature_def({
                "input": receiver_tensor
            }, {"some_output_3": output_3})
    }

    self.assertDictEqual(expected_signature_defs, signature_defs)

  def testGetSignatureDefByKeyClassification(self):
    input1 = constant_op.constant("a", name="input-1")
    output1 = constant_op.constant("b", name="output-1")
    output2 = constant_op.constant(3.0, name="output-2")

    meta_graph_def = meta_graph_pb2.MetaGraphDef()
    self._add_to_signature_def_map(meta_graph_def, {
        "my_classification":
            signature_def_utils.classification_signature_def(
                input1, output1, output2)
    })

    # Look up the classification signature def with the key used while saving.
    signature_def = signature_def_contrib_utils.get_signature_def_by_key(
        meta_graph_def, "my_classification")

    # Check the method name to match the constants classification method name.
    self.assertEqual(signature_constants.CLASSIFY_METHOD_NAME,
                     signature_def.method_name)

    # Check inputs in signature def.
    self.assertEqual(1, len(signature_def.inputs))
    self._check_tensor_info(signature_def.inputs,
                            signature_constants.CLASSIFY_INPUTS, "input-1:0")

    # Check outputs in signature def.
    self.assertEqual(2, len(signature_def.outputs))
    self._check_tensor_info(signature_def.outputs,
                            signature_constants.CLASSIFY_OUTPUT_CLASSES,
                            "output-1:0")
    self._check_tensor_info(signature_def.outputs,
                            signature_constants.CLASSIFY_OUTPUT_SCORES,
                            "output-2:0")

  def _writeDummySavedModel(self, path, feature_name):
    """Writes a classifier with two input features to the given path."""
    with ops.Graph().as_default():
      examples = array_ops.placeholder(dtypes.string, name="input_node")
      feature_configs = {
          feature_name: parsing_ops.FixedLenFeature(shape=[],
                                                    dtype=dtypes.float32),
      }
      features = parsing_ops.parse_example(examples, feature_configs)
      feature = features[feature_name]

      variable_node = variables.VariableV1(1.0, name="variable_node")
      scores = math_ops.multiply(variable_node, feature, name="output_node")
      class_feature = array_ops.fill(array_ops.shape(feature),
                                     "class_%s" % feature_name)
      classes = array_ops.transpose(class_feature)

      with session.Session() as sess:
        sess.run(variables.global_variables_initializer())
        signature = (
            signature_def_utils.classification_signature_def(
                examples=examples,
                classes=classes,
                scores=scores,))
        builder = saved_model_builder.SavedModelBuilder(path)
        builder.add_meta_graph_and_variables(
            sess,
            [tag_constants.SERVING],
            signature_def_map={
                signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                    signature,
            },)
        builder.save(as_text=True)

  def testClassificationSignatureDef(self):
    input1 = constant_op.constant("a", name="input-1")
    output1 = constant_op.constant("b", name="output-1")
    output2 = constant_op.constant("c", name="output-2")
    signature_def = signature_def_utils.classification_signature_def(input1,
                                                                     output1,
                                                                     output2)

    self.assertEqual(signature_constants.CLASSIFY_METHOD_NAME,
                     signature_def.method_name)

    # Check inputs in signature def.
    self.assertEqual(1, len(signature_def.inputs))
    x_tensor_info_actual = (
        signature_def.inputs[signature_constants.CLASSIFY_INPUTS])
    self.assertEqual("input-1:0", x_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, x_tensor_info_actual.dtype)
    self.assertEqual(0, len(x_tensor_info_actual.tensor_shape.dim))

    # Check outputs in signature def.
    self.assertEqual(2, len(signature_def.outputs))
    classes_tensor_info_actual = (
        signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_CLASSES])
    self.assertEqual("output-1:0", classes_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, classes_tensor_info_actual.dtype)
    self.assertEqual(0, len(classes_tensor_info_actual.tensor_shape.dim))
    scores_tensor_info_actual = (
        signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_SCORES])
    self.assertEqual("output-2:0", scores_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, scores_tensor_info_actual.dtype)
    self.assertEqual(0, len(scores_tensor_info_actual.tensor_shape.dim))

 def as_signature_def(self, receiver_tensors):
   if len(receiver_tensors) != 1:
     raise ValueError('Classification input must be a single string Tensor; '
                      'got {}'.format(receiver_tensors))
   (_, examples), = receiver_tensors.items()
   if dtypes.as_dtype(examples.dtype) != dtypes.string:
     raise ValueError('Classification input must be a single string Tensor; '
                      'got {}'.format(receiver_tensors))
   return signature_def_utils.classification_signature_def(
       examples, self.classes, self.scores)

  def test_build_all_signature_defs_with_single_alternatives(self):
    receiver_tensor = array_ops.placeholder(dtypes.string)
    receiver_tensors_alternative_1 = array_ops.placeholder(dtypes.int64)
    receiver_tensors_alternative_2 = array_ops.sparse_placeholder(
        dtypes.float32)
    # Note we are passing single Tensors as values of
    # receiver_tensors_alternatives, where normally that is a dict.
    # In this case a dict will be created using the default receiver tensor
    # name "input".
    receiver_tensors_alternatives = {"other1": receiver_tensors_alternative_1,
                                     "other2": receiver_tensors_alternative_2}
    output_1 = constant_op.constant([1.])
    output_2 = constant_op.constant(["2"])
    output_3 = constant_op.constant(["3"])
    export_outputs = {
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
            export_output.RegressionOutput(value=output_1),
        "head-2": export_output.ClassificationOutput(classes=output_2),
        "head-3": export_output.PredictOutput(outputs={
            "some_output_3": output_3
        }),
    }

    signature_defs = export.build_all_signature_defs(
        receiver_tensor, export_outputs, receiver_tensors_alternatives)

    expected_signature_defs = {
        "serving_default":
            signature_def_utils.regression_signature_def(
                receiver_tensor,
                output_1),
        "head-2":
            signature_def_utils.classification_signature_def(
                receiver_tensor,
                output_2, None),
        "head-3":
            signature_def_utils.predict_signature_def(
                {"input": receiver_tensor},
                {"some_output_3": output_3}),
        "other1:head-3":
            signature_def_utils.predict_signature_def(
                {"input": receiver_tensors_alternative_1},
                {"some_output_3": output_3}),
        "other2:head-3":
            signature_def_utils.predict_signature_def(
                {"input": receiver_tensors_alternative_2},
                {"some_output_3": output_3})

        # Note that the alternatives 'other:serving_default' and 'other:head-2'
        # are invalid, because regession and classification signatures must take
        # a single string input.  Here we verify that these invalid signatures
        # are not included in the export.
    }

    self.assertDictEqual(expected_signature_defs, signature_defs)

def build_standardized_signature_def(input_tensors, output_tensors,
                                     problem_type):
  """Build a SignatureDef using problem type and input and output Tensors.

  Note that this delegates the actual creation of the signatures to methods in
  //third_party/tensorflow/python/saved_model/signature_def_utils.py, which may
  assign names to the input and output tensors (depending on the problem type)
  that are standardized in the context of SavedModel.

  Args:
    input_tensors: a dict of string key to `Tensor`
    output_tensors: a dict of string key to `Tensor`
    problem_type: an instance of constants.ProblemType, specifying
      classification, regression, etc.

  Returns:
    A SignatureDef using SavedModel standard keys where possible.

  Raises:
    ValueError: if input_tensors or output_tensors is None or empty.
  """

  if not input_tensors:
    raise ValueError('input_tensors must be provided.')
  if not output_tensors:
    raise ValueError('output_tensors must be provided.')

  # Per-method signature_def functions will standardize the keys if possible
  if _is_classification_problem(problem_type, input_tensors, output_tensors):
    (_, examples), = input_tensors.items()
    classes = _get_classification_classes(output_tensors)
    scores = _get_classification_scores(output_tensors)
    if classes is None and scores is None:
      items = list(output_tensors.items())
      if items[0][1].dtype == dtypes.string:
        (_, classes), = items
      else:
        (_, scores), = items
    return signature_def_utils.classification_signature_def(
        examples, classes, scores)
  elif _is_regression_problem(problem_type, input_tensors, output_tensors):
    (_, examples), = input_tensors.items()
    (_, predictions), = output_tensors.items()
    return signature_def_utils.regression_signature_def(examples, predictions)
  else:
    return signature_def_utils.predict_signature_def(input_tensors,
                                                     output_tensors)

  def test_build_all_signature_defs_with_dict_alternatives(self):
    receiver_tensor = array_ops.placeholder(dtypes.string)
    receiver_tensors_alternative_1 = {
        "foo": array_ops.placeholder(dtypes.int64),
        "bar": array_ops.sparse_placeholder(dtypes.float32)}
    receiver_tensors_alternatives = {"other": receiver_tensors_alternative_1}
    output_1 = constant_op.constant([1.])
    output_2 = constant_op.constant(["2"])
    output_3 = constant_op.constant(["3"])
    export_outputs = {
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
            export_output.RegressionOutput(value=output_1),
        "head-2": export_output.ClassificationOutput(classes=output_2),
        "head-3": export_output.PredictOutput(outputs={
            "some_output_3": output_3
        }),
    }

    signature_defs = export_utils.build_all_signature_defs(
        receiver_tensor, export_outputs, receiver_tensors_alternatives)

    expected_signature_defs = {
        "serving_default":
            signature_def_utils.regression_signature_def(
                receiver_tensor,
                output_1),
        "head-2":
            signature_def_utils.classification_signature_def(
                receiver_tensor,
                output_2, None),
        "head-3":
            signature_def_utils.predict_signature_def(
                {"input": receiver_tensor},
                {"some_output_3": output_3}),
        "other:head-3":
            signature_def_utils.predict_signature_def(
                receiver_tensors_alternative_1,
                {"some_output_3": output_3})

        # Note that the alternatives 'other:serving_default' and
        # 'other:head-2' are invalid, because regession and classification
        # signatures must take a single string input.  Here we verify that
        # these invalid signatures are not included in the export_utils.
    }

    self.assertDictEqual(expected_signature_defs, signature_defs)