make predictions respect the event_level

Author: hfrick

R/boost_tree.R

@@ -472,6 +472,17 @@ as_xgb_data <- function(x, y, validation = 0, event_level = "first", ...) {
 
   list(data = dat, watchlist = wlist)
 }
+
+get_event_level <- function(model_spec){
+  if ("event_level" %in% names(model_spec$eng_args)) {
+    event_level <- get_expr(model_spec$eng_args$event_level)
+  } else {
+    # "first" is the default for as_xgb_data() and xgb_train()
+    event_level <- "first"
+  }
+  event_level
+}
+
 #' @importFrom purrr map_df
 #' @export
 #' @rdname multi_predict

R/boost_tree_data.R

@@ -158,7 +158,12 @@ set_pred(
     pre = NULL,
     post = function(x, object) {
       if (is.vector(x)) {
-        x <- ifelse(x >= 0.5, object$lvl[2], object$lvl[1])
+        event_level <- get_event_level(object$spec)
+        if (event_level == "first") {
+          x <- ifelse(x >= 0.5, object$lvl[1], object$lvl[2])
+        } else {
+          x <- ifelse(x >= 0.5, object$lvl[2], object$lvl[1])
+        }
       } else {
         x <- object$lvl[apply(x, 1, which.max)]
       }
@@ -178,7 +183,12 @@ set_pred(
     pre = NULL,
     post = function(x, object) {
       if (is.vector(x)) {
-        x <- tibble(v1 = 1 - x, v2 = x)
+        event_level <- get_event_level(object$spec)
+        if (event_level == "first") {
+          x <- tibble(v1 = x, v2 = 1 - x)
+        } else {
+          x <- tibble(v1 = 1 - x, v2 = x)
+        }
       } else {
         x <- as_tibble(x, .name_repair = "minimal")
       }