Converts a PredictResponse to ClassificationResponse or RegressionResponse.

  Args:
    pred: PredictResponse to convert.
    serving_bundle: A `ServingBundle` object that contains the information about
      the serving request that the response was generated by.

  Returns:
    A ClassificationResponse or RegressionResponse.

def convert_predict_response(pred, serving_bundle):
  """Converts a PredictResponse to ClassificationResponse or RegressionResponse.

  Args:
    pred: PredictResponse to convert.
    serving_bundle: A `ServingBundle` object that contains the information about
      the serving request that the response was generated by.

  Returns:
    A ClassificationResponse or RegressionResponse.
  """
  output = pred.outputs[serving_bundle.predict_output_tensor]
  raw_output = output.float_val
  if serving_bundle.model_type == 'classification':
    values = []
    for example_index in range(output.tensor_shape.dim[0].size):
      start = example_index * output.tensor_shape.dim[1].size
      values.append(raw_output[start:start + output.tensor_shape.dim[1].size])
  else:
    values = raw_output
  return convert_prediction_values(values, serving_bundle, pred.model_spec)