R/kerasnip_spec_methods.R

In kerasnip: A Bridge Between 'keras' and 'tidymodels'

#' set_args Method for kerasnip Spec Objects
#'
#' @description
#' S3 method for `set_args()` dispatched on `kerasnip_spec` objects.
#' `parsnip::set_args.model_spec()` calls `new_model_spec()`, which strips any
#' extra classes and attributes. This wrapper saves and re-attaches the
#' `kerasnip_layer_blocks` and `kerasnip_functional` metadata attributes (and
#' the `kerasnip_spec` class) after `NextMethod()` has done its work.
#'
#' @param object A `kerasnip_spec` model specification.
#' @param ... Named model arguments to update, passed to
#'   `parsnip::set_args.model_spec()`.
#' @return A `model_spec` object with the `kerasnip_spec` class and metadata
#'   attributes re-attached.
#' @keywords internal
#' @importFrom parsnip set_args
#' @exportS3Method parsnip::set_args
set_args.kerasnip_spec <- function(object, ...) {
  layer_blocks <- attr(object, "kerasnip_layer_blocks")
  functional <- attr(object, "kerasnip_functional")
  result <- NextMethod()
  class(result) <- c(class(result)[1L], "kerasnip_spec", class(result)[-1L])
  attr(result, "kerasnip_layer_blocks") <- layer_blocks
  attr(result, "kerasnip_functional") <- functional
  result
}

#' set_engine Method for kerasnip Spec Objects
#'
#' @description
#' S3 method for `set_engine()` dispatched on `kerasnip_spec` objects.
#' `parsnip::set_engine.model_spec()` internally calls `new_model_spec()`,
#' which re-creates the spec from scratch with only `c(cls, "model_spec")` as
#' the class vector — stripping `kerasnip_spec` and any custom attributes. This
#' wrapper preserves the `kerasnip_layer_blocks` and `kerasnip_functional`
#' metadata attributes and re-attaches them (along with the `kerasnip_spec`
#' class) after `NextMethod()` has done its work.
#'
#' @param object A `kerasnip_spec` model specification.
#' @param engine A character string naming the engine (e.g., `"keras"`).
#' @param ... Additional engine-specific arguments passed to
#'   `parsnip::set_engine.model_spec()`.
#' @return A `model_spec` object with the `kerasnip_spec` class and metadata
#'   attributes re-attached.
#' @keywords internal
#' @importFrom parsnip set_engine
#' @exportS3Method parsnip::set_engine
set_engine.kerasnip_spec <- function(object, engine, ...) {
  layer_blocks <- attr(object, "kerasnip_layer_blocks")
  functional <- attr(object, "kerasnip_functional")
  result <- NextMethod()
  class(result) <- c(class(result)[1L], "kerasnip_spec", class(result)[-1L])
  attr(result, "kerasnip_layer_blocks") <- layer_blocks
  attr(result, "kerasnip_functional") <- functional
  result
}

#' Fit Method for kerasnip Spec Objects
#'
#' @description
#' S3 method for `fit()` dispatched on `kerasnip_spec` objects. Delegates to
#' the standard parsnip `fit.model_spec()` and then tags the result with the
#' `kerasnip_model_fit` class so that `predict.kerasnip_model_fit()` is
#' dispatched on subsequent calls.
#'
#' `kerasnip_spec` is stripped from the class before `NextMethod()` to prevent
#' parsnip's internal `specific_model()` helper from returning more than one
#' model-class entry, which would break registry lookups. The custom metadata
#' attributes remain on the object and are thus stored inside the resulting
#' `model_fit$spec`.
#'
#' @param object A `kerasnip_spec` model specification.
#' @param ... Passed to `parsnip::fit.model_spec()`.
#' @return A `model_fit` object with the additional `kerasnip_model_fit` class
#'   prepended to its class vector.
#' @keywords internal
#' @importFrom generics fit
#' @exportS3Method generics::fit
fit.kerasnip_spec <- function(object, ...) {
  class(object) <- class(object)[class(object) != "kerasnip_spec"]
  result <- NextMethod()
  class(result) <- c("kerasnip_model_fit", class(result))
  result
}

#' fit_xy Method for kerasnip Spec Objects
#'
#' @description
#' S3 method for `fit_xy()` dispatched on `kerasnip_spec` objects. Workflows
#' route through `fit_xy` rather than `fit`, so this method ensures the
#' `kerasnip_model_fit` class is attached in the workflow fitting path as well.
#'
#' `kerasnip_spec` is stripped from the class before `NextMethod()` to prevent
#' parsnip's internal `specific_model()` helper from returning more than one
#' model-class entry, which would break registry lookups. The custom metadata
#' attributes remain on the object and are thus stored inside the resulting
#' `model_fit$spec`.
#'
#' @param object A `kerasnip_spec` model specification.
#' @param ... Passed to `parsnip::fit_xy.model_spec()`.
#' @return A `model_fit` object with the additional `kerasnip_model_fit` class
#'   prepended to its class vector.
#' @keywords internal
#' @importFrom generics fit_xy
#' @exportS3Method generics::fit_xy
fit_xy.kerasnip_spec <- function(object, ...) {
  class(object) <- class(object)[class(object) != "kerasnip_spec"]
  result <- NextMethod()
  class(result) <- c("kerasnip_model_fit", class(result))
  result
}

#' Predict Method for kerasnip Model Fits
#'
#' @description
#' S3 method for `predict()` dispatched on `kerasnip_model_fit` objects.
#' Before delegating to the standard parsnip predict machinery, it checks
#' whether the underlying model type is registered in the current parsnip
#' session. If not (e.g. after loading a saved workflow in a new R session),
#' it transparently replays the full parsnip registration using metadata stored
#' on the spec object — requiring no manual step from the user.
#'
#' @details
#' The metadata needed for re-registration (`kerasnip_layer_blocks`,
#' `kerasnip_functional`) is embedded on the spec object by the spec
#' constructor function at call time. This means it is preserved across
#' `saveRDS()`/`readRDS()` and `bundle()`/`unbundle()` round-trips.
#'
#' For full model weight portability (i.e. to be able to `predict()` on new
#' data in a new R session), use `bundle::bundle()` before saving. Plain
#' `saveRDS()` preserves the spec structure and will auto-register, but the
#' underlying Keras model weights are not portable without bundling.
#'
#' @param object A `kerasnip_model_fit` object.
#' @param new_data A data frame of predictors.
#' @param ... Passed to the parsnip predict method.
#' @return A tibble of predictions.
#' @keywords internal
#' @importFrom stats predict
#' @exportS3Method stats::predict
predict.kerasnip_model_fit <- function(object, new_data, ...) {
  model_name <- class(object$spec)[1L]

  if (!model_exists(model_name)) {
    spec <- object$spec
    layer_blocks <- attr(spec, "kerasnip_layer_blocks")
    functional <- attr(spec, "kerasnip_functional") %||% FALSE
    mode <- spec$mode
    args_info <- collect_spec_args(layer_blocks, functional)

    register_core_model(model_name, mode)
    register_model_args(model_name, args_info$parsnip_names)
    register_fit_predict(model_name, mode, layer_blocks, functional)
    register_update_method(
      model_name,
      args_info$parsnip_names,
      env = globalenv()
    )
  }

  # Restore Keras model from serialized bytes if the Python object has been
  # invalidated (e.g. by saveRDS()/readRDS() within the same session, or after
  # bundle()/unbundle()). reticulate::py_validate_xptr() performs a direct
  # C-level check and throws the same "Unable to access object" error that
  # would otherwise surface later during predict(); catching it here lets us
  # restore the model transparently before dispatching to parsnip.
  if (!is.null(object$fit$keras_bytes)) {
    is_valid <- tryCatch(
      {
        reticulate::py_validate_xptr(object$fit$fit)
        TRUE
      },
      error = function(e) FALSE
    )
    if (!is_valid) {
      object$fit$fit <- keras_model_from_bytes(object$fit$keras_bytes)
    }
  }

  # Restore Laplace combined models from serialized bytes if present
  if (!is.null(object$fit$laplace)) {
    for (nm in names(object$fit$laplace)) {
      entry <- object$fit$laplace[[nm]]
      if (
        !is.null(entry$combined_model_bytes) && !is.null(entry$combined_model)
      ) {
        is_valid <- tryCatch(
          {
            reticulate::py_validate_xptr(entry$combined_model)
            TRUE
          },
          error = function(e) FALSE
        )
        if (!is_valid) {
          object$fit$laplace[[nm]]$combined_model <-
            keras_model_from_bytes(entry$combined_model_bytes)
        }
      }
    }
  }

  # Strip our class and dispatch explicitly. NextMethod() re-uses the original
  # call arguments, not locally modified ones, so the restored model above
  # would not be forwarded via NextMethod().
  class(object) <- class(object)[class(object) != "kerasnip_model_fit"]

  dots <- list(...)
  type <- dots$type %||%
    if (object$spec$mode == "classification") "class" else "numeric"
  is_multi_output_class <- object$spec$mode == "classification" &&
    is.list(object$fit$lvl) &&
    !is.null(names(object$fit$lvl))

  if (type == "class" && is_multi_output_class) {
    # parsnip::predict_class.model_fit() assumes `type = "class"` results are
    # either a bare factor/vector or a single-column data frame; for any
    # result with more than one column it does `res$values <- factor(...,
    # levels = object$lvl)`, which errors because our multi-output tibble has
    # no `values` column and `object$lvl` (parsnip's own top-level field, not
    # `object$fit$lvl`) isn't meaningful for multiple, independently-leveled
    # outputs. Bypass that tail end and call our own registered pred
    # pipeline directly; `type = "prob"` doesn't have this issue and needs no
    # such bypass.
    return(predict_class_multi_output(object, new_data))
  }

  # `joint = TRUE` is a kerasnip-specific argument, not part of parsnip's
  # `predict()` API. parsnip's `type` argument is validated against a
  # fixed, hardcoded list and its dispatch has no extension point, so this
  # can't be registered as a new `type`. It's intercepted here instead, and
  # stripped from `dots` before any fall-through to standard dispatch so it
  # never reaches parsnip's own predict machinery.
  joint <- isTRUE(dots$joint)
  n_draws <- dots$n_draws %||% 1000L
  dots[c("joint", "n_draws")] <- NULL

  if (joint) {
    if (!identical(type, "pred_int")) {
      rlang::abort(c(
        "`joint = TRUE` is only supported for `type = \"pred_int\"`.",
        i = "Confidence intervals reflect epistemic (weight) uncertainty",
        i = "only, and this implementation has no estimated source of",
        i = "cross-step correlation for that case.",
        i = paste0("Got `type = \"", type, "\"`.")
      ))
    }
    return(laplace_joint_pred_int(object, new_data, n_draws = n_draws))
  }

  rlang::exec(predict, object, new_data = new_data, !!!dots)
}

# Replicates parsnip::predict_class.model_fit() up to (but not including) its
# final single-output factor-coercion step, which is incompatible with
# multi-output classification results. See predict.kerasnip_model_fit().
#
# `new_data` is used as-is rather than routed through parsnip's internal
# prepare_data(): unnecessary here since kerasnip's own registered pred
# pipeline (process_x_functional()/process_x_sequential(), invoked via
# pred_class$args) operates on the full data frame directly and doesn't rely
# on prepare_data()'s x_names column-subsetting. Likewise, `pred_class$pre`
# and a namespaced `pred_class$func` are always NULL/absent for kerasnip's
# own registered `type = "class"` entry (register_fit_predict.R), so unlike
# parsnip:::make_pred_call() this doesn't need to handle either case.
predict_class_multi_output <- function(object, new_data) {
  pred_class <- object$spec$method$pred$class
  pred_call <- rlang::call2(pred_class$func["fun"], !!!pred_class$args)
  res <- rlang::eval_tidy(pred_call)
  if (!is.null(pred_class$post)) {
    res <- pred_class$post(res, object)
  }
  res
}