def model_fn()

in src/entrypoint/inference.py [0:0]

• 14 lines of code
• 3 McCabe index (conditional complexity)

def model_fn(model_dir: Union[str, Path]) -> Predictor:
    """Load a glounts model from a directory.

    Args:
        model_dir (Union[str, Path]): a directory where model is saved.

    Returns:
        Predictor: A gluonts predictor.
    """
    predictor = Predictor.deserialize(Path(model_dir))

    # If model was trained on log-space, then forecast must be inverted before metrics etc.
    with open(os.path.join(model_dir, "y_transform.json"), "r") as f:
        y_transform = json.load(f)
        logger.info("model_fn: custom transformations = %s", y_transform)

        if y_transform["inverse_transform"] == "expm1":
            predictor.output_transform = expm1_and_clip_to_zero
        else:
            predictor.output_transform = clip_to_zero

        # Custom field
        predictor.pre_input_transform = log1p if y_transform["transform"] == "log1p" else None

    logger.info("predictor.pre_input_transform: %s", predictor.pre_input_transform)
    logger.info("predictor.output_transform: %s", predictor.output_transform)
    logger.info("model_fn() done; loaded predictor %s", predictor)

    return predictor