def default_output_fn()

in src/sagemaker_xgboost_container/algorithm_mode/handler_service.py [0:0]

• 21 lines of code
• 9 McCabe index (conditional complexity)

        def default_output_fn(self, prediction, accept):
            """Return encoded prediction for the response.
            Args:
                prediction (obj): prediction returned by predict_fn .
                accept (str): accept content-type expected by the client.
            Returns:
                encoded response for MMS to return to client
            """
            accept_type = accept.lower()
            try:
                if accept_type == content_types.CSV or accept_type == "csv":
                    if SAGEMAKER_BATCH:
                        return_data = "\n".join(map(str, prediction.tolist())) + "\n"
                    else:
                        # FIXME: this is invalid CSV and is only retained for backwards compatibility
                        return_data = ",".join(map(str, prediction.tolist()))
                    encoded_prediction = return_data.encode("utf-8")
                elif accept_type == content_types.JSON or accept_type == "json":
                    encoded_prediction = encoder.encode(prediction, accept_type)
                else:
                    raise ValueError(
                        "{} is not an accepted Accept type. Please choose one of the following:"
                        " ['{}', '{}'].".format(accept, content_types.CSV, content_types.JSON)
                    )
            except Exception as e:
                raise UnsupportedMediaTypeInferenceError(
                    "Encoding to accept type {} failed with exception: {}".format(accept, e)
                )
            return encoded_prediction