Merge pull request #268 from tidymodels/using-parsnip-models

adjust for cases where parsnip predictions are from a parsnip model

Author: topepo

NEWS.md

@@ -1,5 +1,9 @@
 # parsnip (development version)
 
+# parsnip 0.0.5.9000
+
+
+
 # parsnip 0.0.5
 
 ## Fixes

R/predict.R

@@ -144,7 +144,6 @@ predict.model_fit <- function(object, new_data, type = NULL, opts = list(), ...)
     raw      = predict_raw(object = object, new_data = new_data, opts = opts, ...),
     rlang::abort(glue::glue("I don't know about type = '{type}'"))
   )
-
   if (!inherits(res, "tbl_spark")) {
     res <- switch(
       type,
@@ -186,9 +185,11 @@ format_num <- function(x) {
   if (inherits(x, "tbl_spark"))
     return(x)
 
-  if (isTRUE(ncol(x) > 1)) {
+  if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
     x <- as_tibble(x, .name_repair = "minimal")
-    names(x) <- paste0(".pred_", names(x))
+    if (!any(grepl("^\\.pred", names(x)))) {
+      names(x) <- paste0(".pred_", names(x))
+    }
   } else {
     x <- tibble(.pred = x)
   }
@@ -204,7 +205,9 @@ format_class <- function(x) {
 }
 
 format_classprobs <- function(x) {
-  names(x) <- paste0(".pred_", names(x))
+  if (!any(grepl("^\\.pred_", names(x)))) {
+    names(x) <- paste0(".pred_", names(x))
+  }
   x <- as_tibble(x)
   x
 }

R/predict_class.R

@@ -40,8 +40,14 @@ predict_class.model_fit <- function(object, new_data, ...) {
   if (is.vector(res) || is.factor(res)) {
     res <- factor(as.character(res), levels = object$lvl)
   } else {
-    if (!inherits(res, "tbl_spark"))
-      res$values <- factor(as.character(res$values), levels = object$lvl)
+    if (!inherits(res, "tbl_spark")) {
+      # Now case where a parsnip model generated `res`
+      if (is.data.frame(res) && ncol(res) == 1 && is.factor(res[[1]])) {
+        res <- res[[1]]
+      } else {
+        res$values <- factor(as.character(res$values), levels = object$lvl)
+      }
+    }
   }
 
   res