add the mlverse:tabnet model familly

DESCRIPTION

@@ -62,6 +62,7 @@ Suggests:
     rsample,
     RSpectra,
     survival (>= 3.2-10),
+    tabnet,
     testthat (>= 3.0.0),
     TH.data,
     usethis (>= 1.5.0),

NAMESPACE

@@ -86,6 +86,7 @@ S3method(axe_env,terms)
 S3method(axe_env,train)
 S3method(axe_env,train.recipe)
 S3method(axe_env,xgb.Booster)
+S3method(axe_fitted,"_tabnet_fit")
 S3method(axe_fitted,C5.0)
 S3method(axe_fitted,default)
 S3method(axe_fitted,earth)

NEWS.md

@@ -1,5 +1,7 @@
 # butcher (development version)
 
+* Added butcher methods for `tabnet()` (@cregouby #226).
+
 # butcher 0.3.0
 
 * Julia Silge is now the maintainer (#230).

R/tabnet_fit.R

@@ -0,0 +1,45 @@
+#' Axing a tabnet_fit.
+#'
+#' @inheritParams butcher
+#'
+#' @return Axed tabnet_fit object.
+#'
+#' @examplesIf rlang::is_installed("tabnet")
+#'
+#' # Load libraries
+#' suppressWarnings(suppressMessages(library(parsnip)))
+#' suppressWarnings(suppressMessages(library(rsample)))
+#'
+#' # Load data
+#' split <- initial_split(mtcars, props = 9/10)
+#' car_train <- training(split)
+#'
+#' # Create model and fit
+#' mtcar_fit <- tabnet() %>%
+#'   set_mode("regression") %>%
+#'   set_engine("torch")
+#'   fit(mpg ~ ., data = car_train)
+#'
+#' out <- butcher(mtcar_fit, verbose = TRUE)
+#'
+#' @name axe-tabnet_fit
+NULL
+
+#' Remove fitted values.
+#'
+#' @rdname axe-tabnet_fit
+#' @export
+axe_fitted._tabnet_fit <- function(x, verbose = FALSE, ...) {
+  old <- x
+  x$fit$fit <- exchange(x$fit$fit, "checkpoints", list(NULL))
+  x$fit$fit$importances <- exchange(x$fit$fit$importances, "variables", list(NULL))
+  x$fit$fit$importances <- exchange(x$fit$fit$importances, "importance", list(NULL))
+
+  add_butcher_attributes(
+    x,
+    old,
+    disabled = NULL,
+    add_class = FALSE,
+    verbose = verbose
+  )
+}

tests/testthat/test-tabnet_fit.R

@@ -0,0 +1,41 @@
+test_that("tabnet_fit + axe_fitted() works", {
+  skip_on_cran()
+  skip_if_not_installed("tabnet")
+  suppressPackageStartupMessages(library(parsnip))
+  # Create model and fit
+  tabnet_fit <- tabnet::tabnet(epochs = 10) %>%
+    set_mode("regression") %>%
+    set_engine("torch") %>%
+    fit(mpg ~ ., data = mtcars)
+
+  expect_error(axed_out <- axe_fitted(tabnet_fit, verbose = TRUE), NA)
+  expect_lt(lobstr::obj_size(axed_out),lobstr::obj_size(tabnet_fit))
+})
+
+test_that("tabnet_fit + butcher() works", {
+  skip_on_cran()
+  skip_if_not_installed("tabnet")
+  suppressPackageStartupMessages(library(parsnip))
+  # Create model and fit
+  tabnet_fit <- tabnet::tabnet(epochs = 10) %>%
+    set_mode("regression") %>%
+    set_engine("torch") %>%
+    fit(mpg ~ ., data = mtcars)
+
+  expect_error(tabnet_out <- butcher(tabnet_fit, verbose = TRUE), NA)
+})
+
+test_that("tabnet_fit + predict() works", {
+  skip_on_cran()
+  skip_if_not_installed("tabnet")
+  suppressPackageStartupMessages(library(parsnip))
+  # Create model and fit
+  tabnet_fit <- tabnet::tabnet(epochs = 10) %>%
+    set_mode("regression") %>%
+    set_engine("torch") %>%
+    fit(mpg ~ ., data = mtcars)
+
+  tabnet_out <- butcher(tabnet_fit, verbose = TRUE)
+  new_data <- as.matrix(mtcars[1:3, 2:11])
+  expect_equal(predict(tabnet_out,new_data), predict(tabnet_fit, new_data))
+})