def run_inference(examples, serving_bundle): """Run inference on examples given model information Args: examples: A list of examples that matches the model spec. serving_bundle: A `ServingBundle` object that contains the information to make the inference request. Returns: A tuple with the first entry being the ClassificationResponse or RegressionResponse proto and the second entry being a dictionary of extra data for each example, such as attributions, or None if no data exists. """ batch_size = 64 if serving_bundle.estimator and serving_bundle.feature_spec: # If provided an estimator and feature spec then run inference locally. preds = serving_bundle.estimator.predict( lambda: tf.data.Dataset.from_tensor_slices( tf.io.parse_example([ ex.SerializeToString() for ex in examples ], serving_bundle.feature_spec)).batch(batch_size)) if serving_bundle.use_predict: preds_key = serving_bundle.predict_output_tensor elif serving_bundle.model_type == 'regression': preds_key = 'predictions' else: preds_key = 'probabilities' values = [] for pred in preds: values.append(pred[preds_key]) return (common_utils.convert_prediction_values(values, serving_bundle), None) elif serving_bundle.custom_predict_fn: # If custom_predict_fn is provided, pass examples directly for local # inference. values = serving_bundle.custom_predict_fn(examples) extra_results = None # If the custom prediction function returned a dict, then parse out the # prediction scores. If it is just a list, then the results are the # prediction results without attributions or other data. if isinstance(values, dict): preds = values.pop('predictions') extra_results = values else: preds = values return (common_utils.convert_prediction_values(preds, serving_bundle), extra_results) else: return (platform_utils.call_servo(examples, serving_bundle), None)
def run_inference(examples, serving_bundle): """Run inference on examples given model information Args: examples: A list of examples that matches the model spec. serving_bundle: A `ServingBundle` object that contains the information to make the inference request. Returns: A ClassificationResponse or RegressionResponse proto. """ batch_size = 64 if serving_bundle.estimator and serving_bundle.feature_spec: # If provided an estimator and feature spec then run inference locally. preds = serving_bundle.estimator.predict( lambda: tf.data.Dataset.from_tensor_slices( tf.parse_example([ex.SerializeToString() for ex in examples], serving_bundle. feature_spec)).batch(batch_size)) if serving_bundle.use_predict: preds_key = serving_bundle.predict_output_tensor elif serving_bundle.model_type == 'regression': preds_key = 'predictions' else: preds_key = 'probabilities' values = [] for pred in preds: values.append(pred[preds_key]) return common_utils.convert_prediction_values(values, serving_bundle) elif serving_bundle.custom_predict_fn: # If custom_predict_fn is provided, pass examples directly for local # inference. values = serving_bundle.custom_predict_fn(examples) return common_utils.convert_prediction_values(values, serving_bundle) else: return platform_utils.call_servo(examples, serving_bundle)
def run_inference(examples, serving_bundle): """Run inference on examples given model information. Args: examples: A list of examples that matches the model spec. serving_bundle: A `ServingBundle` object that contains the information to make the inference request. Returns: A tuple with the first entry being the ClassificationResponse or RegressionResponse proto and the second entry being a dictionary of extra data for each example, such as attributions, or None if no data exists. """ batch_size = 64 if serving_bundle.estimator and serving_bundle.feature_spec: # If provided an estimator and feature spec then run inference locally. preds = serving_bundle.estimator.predict( lambda: tf.data.Dataset.from_tensor_slices( tf.io.parse_example( [ex.SerializeToString() for ex in examples], serving_bundle.feature_spec, ) ).batch(batch_size) ) # Use the specified key if one is provided. key_to_use = ( serving_bundle.predict_output_tensor if serving_bundle.use_predict else None ) values = [] for pred in preds: if key_to_use is None: # If the prediction dictionary only contains one key, use it. returned_keys = list(pred.keys()) if len(returned_keys) == 1: key_to_use = returned_keys[0] # Use default keys if necessary. elif serving_bundle.model_type == "classification": key_to_use = "probabilities" else: key_to_use = "predictions" if key_to_use not in pred: raise KeyError( '"%s" not found in model predictions dictionary' % key_to_use ) values.append(pred[key_to_use]) return ( common_utils.convert_prediction_values(values, serving_bundle), None, ) elif serving_bundle.custom_predict_fn: # If custom_predict_fn is provided, pass examples directly for local # inference. values = serving_bundle.custom_predict_fn(examples) extra_results = None # If the custom prediction function returned a dict, then parse out the # prediction scores. If it is just a list, then the results are the # prediction results without attributions or other data. if isinstance(values, dict): preds = values.pop("predictions") extra_results = values else: preds = values return ( common_utils.convert_prediction_values(preds, serving_bundle), extra_results, ) else: return (platform_utils.call_servo(examples, serving_bundle), None)