Source code for sagemaker.serve.spec.inference_spec

"""Implements class that holds custom load and invoke function of a model"""

from __future__ import absolute_import
import abc




class InferenceSpec(abc.ABC):
    """Abstract base class for holding custom ``load``, ``invoke`` and ``prepare`` functions.

    Provides a skeleton for customization to override the methods
    ``load``, ``invoke`` and ``prepare``.
    """



    @abc.abstractmethod
    def load(self, model_dir: str):
        """Loads the model stored in model_dir and return the model object.

        Args:
            model_dir (str): Path to the directory where the model is stored.
        """




    @abc.abstractmethod
    def invoke(self, input_object: object, model: object):
        """Given model object and input, make inference and return the result.

        Args:
            input_object (object): The input to model
            model (object): The model object
        """




    def preprocess(self, input_data: object):
        """Custom pre-processing function"""




    def postprocess(self, predictions: object):
        """Custom post-processing function"""




    def prepare(self, *args, **kwargs):
        """Custom prepare function"""




    def get_model(self):
        """Return HuggingFace model name for inference spec"""