Feature/tabnet model

NAMESPACE

@@ -11,12 +11,14 @@ S3method(handler_predict,default)
 S3method(handler_predict,glm)
 S3method(handler_predict,lm)
 S3method(handler_predict,ranger)
+S3method(handler_predict,tabnet_fit)
 S3method(handler_predict,train)
 S3method(handler_predict,workflow)
 S3method(handler_predict,xgb.Booster)
 S3method(handler_startup,Learner)
 S3method(handler_startup,default)
 S3method(handler_startup,ranger)
+S3method(handler_startup,tabnet_fit)
 S3method(handler_startup,train)
 S3method(handler_startup,workflow)
 S3method(handler_startup,xgb.Booster)
@@ -32,12 +34,14 @@ S3method(vetiver_create_description,default)
 S3method(vetiver_create_description,glm)
 S3method(vetiver_create_description,lm)
 S3method(vetiver_create_description,ranger)
+S3method(vetiver_create_description,tabnet_fit)
 S3method(vetiver_create_description,train)
 S3method(vetiver_create_description,workflow)
 S3method(vetiver_create_description,xgb.Booster)
 S3method(vetiver_create_meta,Learner)
 S3method(vetiver_create_meta,default)
 S3method(vetiver_create_meta,ranger)
+S3method(vetiver_create_meta,tabnet_fit)
 S3method(vetiver_create_meta,train)
 S3method(vetiver_create_meta,workflow)
 S3method(vetiver_create_meta,xgb.Booster)
@@ -46,13 +50,15 @@ S3method(vetiver_prepare_model,default)
 S3method(vetiver_prepare_model,glm)
 S3method(vetiver_prepare_model,lm)
 S3method(vetiver_prepare_model,ranger)
+S3method(vetiver_prepare_model,tabnet_fit)
 S3method(vetiver_prepare_model,train)
 S3method(vetiver_prepare_model,workflow)
 S3method(vetiver_ptype,Learner)
 S3method(vetiver_ptype,default)
 S3method(vetiver_ptype,glm)
 S3method(vetiver_ptype,lm)
 S3method(vetiver_ptype,ranger)
+S3method(vetiver_ptype,tabnet_fit)
 S3method(vetiver_ptype,train)
 S3method(vetiver_ptype,workflow)
 S3method(vetiver_ptype,xgb.Booster)

NEWS.md

@@ -1,5 +1,7 @@
 # vetiver (development version)
 
+* Added support for tabnet models (#124 @cregouby)
+
 * Trailing slashes are now removed from `vetiver_endpoint()` (#134).
 
 # vetiver 0.1.7

R/tabnet.R

@@ -0,0 +1,54 @@
+#' @rdname vetiver_create_description
+#' @export
+vetiver_create_description.tabnet_fit <- function(model) {
+    paste0(
+        "A tabnet `nn_module` containing ",
+        format(sum(sapply(model$fit$network$parameters, function(x) prod(x$shape))), nsmall = 0, big.mark = ",", scientific = FALSE),
+        " parameters."
+  )
+}
+
+#' @rdname vetiver_create_meta
+#' @export
+vetiver_create_meta.tabnet_fit <- function(model, metadata) {
+    vetiver_meta(metadata, required_pkgs = "tabnet")
+}
+
+#' @rdname vetiver_create_description
+#' @export
+vetiver_prepare_model.tabnet_fit <- function(model) {
+    butcher::butcher(model)
+}
+
+#' @rdname vetiver_create_ptype
+#' @export
+vetiver_ptype.tabnet_fit <- function(model, ...) {
+    rlang::check_dots_used()
+    dots <- list(...)
+    check_ptype_data(dots)
+    ptype <- vctrs::vec_ptype(dots$ptype_data)
+    tibble::as_tibble(ptype)
+}
+
+#' @rdname handler_startup
+#' @export
+handler_startup.tabnet_fit <- function(vetiver_model) {
+    attach_pkgs("tabnet")
+}
+
+#' @rdname handler_startup
+#' @export
+handler_predict.tabnet_fit <- function(vetiver_model, ...) {
+
+    ptype <- vetiver_model$blueprint$ptypes
+
+    function(req) {
+        new_data <- req$body
+        if (!is_null(ptype)) {
+            new_data <- vetiver_type_convert(new_data, ptype)
+            new_data <- hardhat::scream(new_data, ptype)
+        }
+        ret <- predict(vetiver_model, data = new_data, ...)
+        list(.pred = ret$.pred)
+    }
+}

tests/testthat/_snaps/tabnet.md

@@ -0,0 +1,41 @@
+# can print tabnet model
+
+    Code
+      v
+    Output
+
+      -- cars3 - <tabnet_fit> model for deployment 
+      A tabnet `nn_module` containing 6,301 parameters. using 10 features
+
+# error for no ptype_data with tabnet
+
+    Code
+      vetiver_model(car_tn, "cars3")
+    Condition
+      Error in `vetiver_create_description()`:
+      ! object 'car_tn' not found
+
+# create plumber.R for tabnet
+
+    Code
+      cat(readr::read_lines(tmp), sep = "\n")
+    Output
+      # Generated by the vetiver package; edit with care
+
+      library(pins)
+      library(plumber)
+      library(rapidoc)
+      library(vetiver)
+
+      # Packages needed to generate model predictions
+      if (FALSE) {
+          library(tabnet)
+      }
+      b <- board_folder(path = "<redacted>")
+      v <- vetiver_pin_read(b, "cars3")
+
+      #* @plumber
+      function(pr) {
+          pr %>% vetiver_api(v)
+      }
+