output_fn

sklearn_nearest_neighbors.py

from sagemaker_containers.beta.framework import worker, encoders
from six import BytesIO

def _npy_dumps(data):
    # Serializes a numpy array into a stream of npy-formatted bytes.
    buffer = BytesIO()
    np.save(buffer, data)
    return buffer.getvalue()

def output_fn(prediction_output, accept):
    if accept == 'application/x-npy':
        print('output_fn input is', prediction_output, 'in format', accept)
        return _npy_dumps(prediction_output), 'application/x-npy'
    elif accept == 'application/json':
        print('output_fn input is', prediction_output, 'in format', accept)
        return worker.Response(encoders.encode(prediction_output, accept), accept, mimetype=accept)
    else:
        raise ValueError('Accept header must be application/x-npy or application/json, but it is {}'.format(accept))