def get_selected_predictions()

in src/sagemaker_xgboost_container/algorithm_mode/serve_utils.py [0:0]
34 lines of code
13 McCabe index (conditional complexity)

def get_selected_predictions(raw_predictions, selected_keys, objective, num_class=""):
    """Build the selected prediction dictionary based on the objective function and
    requested information.

    'raw_predictions' is the output of ScoringService.predict(...) and will change
    depending on the objective xgboost training function used. For each prediction,
    a new dictionary will be built with the selected content requested in 'selected_keys'.

    VALID_OBJECTIVES contains a mapping of objective functions to valid selected keys.
    For example, a booster trained with a "reg:linear" objective function does not output
    'predicted_label' or 'probabilities' (classification content). Invalid keys will be included
    in the response with an np.nan value.

    :param raw_predictions: output of xgboost predict (list of numpy objects)
    :param selected_keys: strings denoting selected keys (list of str)
    :param objective: objective xgboost training function (str)
    :param num_class: number of classes for multiclass classification (str, optional)
    :return: selected prediction (list of dict)
    """
    if objective not in VALID_OBJECTIVES:
        raise ValueError("Objective `{}` unsupported for selectable inference predictions.".format(objective))

    valid_selected_keys = set(selected_keys).intersection(VALID_OBJECTIVES[objective])
    invalid_selected_keys = set(selected_keys).difference(VALID_OBJECTIVES[objective])
    if invalid_selected_keys:
        logging.warning(
            "Selected key(s) {} incompatible for objective '{}'. "
            "Please use list of compatible selectable inference predictions: {}".format(
                invalid_selected_keys, objective, VALID_OBJECTIVES[objective]
            )
        )

    predictions = []
    for raw_prediction in raw_predictions:
        output = {}
        if PREDICTED_LABEL in valid_selected_keys:
            output[PREDICTED_LABEL] = _get_predicted_label(objective, raw_prediction)
        if LABELS in valid_selected_keys:
            output[LABELS] = _get_labels(objective, num_class=num_class)
        if PROBABILITY in valid_selected_keys:
            output[PROBABILITY] = _get_probability(objective, raw_prediction)
        if PROBABILITIES in valid_selected_keys:
            output[PROBABILITIES] = _get_probabilities(objective, raw_prediction)
        if RAW_SCORE in valid_selected_keys:
            output[RAW_SCORE] = _get_raw_score(objective, raw_prediction)
        if RAW_SCORES in valid_selected_keys:
            output[RAW_SCORES] = _get_raw_scores(objective, raw_prediction)
        if PREDICTED_SCORE in valid_selected_keys:
            output[PREDICTED_SCORE] = raw_prediction.item()
        if invalid_selected_keys:
            for invalid_selected_key in invalid_selected_keys:
                output[invalid_selected_key] = np.nan
        predictions.append(output)
    return predictions