# score.py {{SCORE_VERSION}} generated by ADS {{ADS_VERSION}}
import os
import sys
from functools import lru_cache
from transformers import pipeline
import json
from typing import Dict, List
import numpy as np
from io import BytesIO
import base64
import logging
import torch


task = '{{task}}'

"""
   Inference script. This script is used for prediction by scoring server when schema is known.
"""

@lru_cache(maxsize=10)
def load_model(model_file_name=""):
    """
    Loads model from the serialized format

    Returns
    -------
    model:  a model instance on which predict API can be invoked.

    """
    model_dir = os.path.dirname(os.path.realpath(__file__))
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)

    try:

        model = pipeline(task, model = model_dir)
        print("Model is successfully loaded.")
    except Exception as e:
        raise Exception(f'Failed to load the model. ' + str(e))

    return model


def deserialize(data):
    """
    Deserialize json-serialized data to data in original type when sent to
predict.

    Parameters
    ----------
    data: serialized input data.

    Returns
    -------
    data: deserialized input data.

    """
{% if data_deserializer == "json" %}
    if isinstance(data, bytes):
        logging.warning(
            "bytes are passed directly to the model. If the model expects a specific data format, you need to write the conversion logic in `deserialize()` yourself."
        )
        return data

    data_type = data.get('data_type', '') if isinstance(data, dict) else ''
    json_data = data.get('data', data) if isinstance(data, dict) else data

    if "numpy.ndarray" in data_type:
        load_bytes = BytesIO(base64.b64decode(json_data.encode('utf-8')))
        return np.load(load_bytes, allow_pickle=True)
    if "torch.Tensor" in data_type:
        load_bytes = BytesIO(base64.b64decode(json_data.encode('utf-8')))
        return torch.load(load_bytes)
    if isinstance(json_data, bytes):
        import cloudpickle
        return cloudpickle.loads(json_data)

    return json_data

{% elif data_deserializer == "cloudpickle" %}
    import cloudpickle
    from pickle import UnpicklingError
    deserialized_data = data
    try:
        deserialized_data = cloudpickle.loads(data)
    except TypeError:
        pass
    except UnpicklingError:
        logger.warning(
            "bytes are passed directly to the model. If the model expects a specific data format, you need to write the conversion logic in `deserialize()` yourself."
        )

    return deserialized_data

{% else %}
    # Add further data deserialization if needed
    return data
{% endif %}

def pre_inference(data):
    """
    Preprocess json-serialized data to feed into predict function.

    Parameters
    ----------
    data: Data format as expected by the predict API of the core estimator.

    Returns
    -------
    data: Data format after any processing.

    """
    data = deserialize(data)

    # Add further data preprocessing if needed
    return data

def post_inference(yhat):
    """
    Post-process the model results.

    Parameters
    ----------
    yhat: Data format after calling model.predict.

    Returns
    -------
    yhat: Data format after any processing.

    """
    if isinstance(yhat, dict):
        return serialize_prediction(yhat)
    elif isinstance(yhat, list):
        serialized_yhat = []
        for yhat_i in yhat:
            serialized_yhat.append(serialize_prediction(yhat_i))
        return serialized_yhat
    elif isinstance(yhat, torch.Tensor):
        return yhat.tolist()

    if task == "conversational":
        raise NotImplementedError("The model output of task `conversational` is not json serializable. Add code here to populate json serializable output.")

    # Add further data postprocessing if needed
    return yhat


def serialize_prediction(yhat):
    """
    Serialize the values in the prediction.

    Parameters
    ----------
    yhat: model prediction.

    Returns
    -------
    yhat: Serialized model prediction.

    """
    if isinstance(yhat, dict):
        for key, value in yhat.items():
            if isinstance(value, torch.Tensor):
                yhat[key] = value.tolist()
            elif str(type(value)).startswith("<class 'PIL."):
                import PIL
                yhat[key] = np.asarray(value).tolist() if isinstance(value, PIL.Image.Image) else value
    return yhat


def validated_model_outputs(yhat):
    """
    Validates if the prediction is json serializable.

    Parameters
    ----------
    yhat: Prediction.

    Returns
    -------
    None

    Raises
    ------
    TypeError
        Type is not json serializable.

    """
    try:
        json.dumps(yhat)
    except:
        raise TypeError("The output is not json serializable.")


def predict(data, model=load_model()):
    """
    Returns prediction given the model and data to predict.

    Parameters
    ----------
    model: Model instance returned by load_model API
    data: Data format as expected by the predict API of the core estimator

    Returns
    -------
    predictions: Output from scoring server
        Format: {'prediction': output from model.predict method}

    """
    inputs = pre_inference(data)

    yhat = post_inference(
    model(**inputs) if isinstance(inputs, dict) else model(inputs)
    )
    validated_model_outputs(yhat)
    return {'prediction': yhat}
