Added quantile regression

DESCRIPTION

@@ -21,7 +21,7 @@ Imports:
     cli,
     dplyr,
     generics,
-    hardhat,
+    hardhat (>= 1.4.1),
     jsonlite,
     purrr,
     reticulate (>= 1.41.0.1),

R/TabPFN-predict.R

@@ -7,6 +7,10 @@
 #' @param type The type of prediction. For classification, can be `"class"` or
 #' `"prob"`. Defaults to `NULL` which gives all prediction types possible.
 #'
+#' @param quantile_levels A numeric vector of probabilities, sorted in
+#' increasing order, at which to predict the outcome distribution. Regression
+#' only; defaults to `NULL` for no quantile predictions.
+#'
 #' @param ... Not used, but required for extensibility.
 #'
 #' @return
@@ -20,6 +24,9 @@
 #' probability estimates are in columns with the pattern `.pred_{level}` where
 #' `level` is the levels of the outcome factor vector.
 #'
+#' When `quantile_levels` is given, regression results also have a
+#' `.pred_quantile` column of [hardhat::quantile_pred()] values.
+#'
 #' @examples
 #' # Minimal example for quick execution
 #' car_train <- mtcars[ 1:5,   ]
@@ -32,15 +39,34 @@
 #'
 #'  # Predict
 #'  predict(mod, car_test)
+#'  predict(mod, car_test, quantile_levels = c(0.1, 0.5, 0.9))
 #'  augment(mod, car_test)
 #' }
 #' }
 #'
 #' @export
-predict.tab_pfn <- function(object, new_data, type = NULL, ...) {
+predict.tab_pfn <- function(
+  object,
+  new_data,
+  type = NULL,
+  quantile_levels = NULL,
+  ...
+) {
   rlang::check_dots_empty()
+  if (!is.null(quantile_levels) && !is.null(object$levels)) {
+    cli::cli_abort("{.arg quantile_levels} is only for regression models.")
+  }
+  if (!is.null(quantile_levels)) {
+    hardhat::check_quantile_levels(quantile_levels)
+  }
   forged <- hardhat::forge(new_data, object$blueprint)$predictors
-  res <- predict(object$fit, forged, object$levels, type = type)
+  res <- predict(
+    object$fit,
+    forged,
+    object$levels,
+    type = type,
+    quantile_levels = unname(quantile_levels)
+  )
   res
 }
 
@@ -53,17 +79,33 @@ predict.tabpfn.regressor.TabPFNRegressor <- function(
   new_data,
   levels,
   type = NULL,
+  quantile_levels = NULL,
   ...
 ) {
   py_msg <- reticulate::py_capture_output(
-    res <- try(object$predict(new_data), silent = TRUE)
+    res <- try(
+      object$predict(
+        new_data,
+        output_type = if (is.null(quantile_levels)) "mean" else "main",
+        quantiles = as.list(quantile_levels)
+      ),
+      silent = TRUE
+    )
   )
 
   if (inherits(res, "try-error")) {
     msgs <- as.character(res)
     cli::cli_abort("Prediction failed: {msgs}")
-  } else {
+  } else if (is.null(quantile_levels)) {
     res <- tibble::tibble(.pred = as.vector(res))
+  } else {
+    res <- tibble::tibble(
+      .pred = as.vector(res$mean),
+      .pred_quantile = hardhat::quantile_pred(
+        do.call(cbind, res$quantiles),
+        quantile_levels
+      )
+    )
   }
 
   res
@@ -106,9 +148,15 @@ predict.tabpfn.classifier.TabPFNClassifier <- function(
 
 #' @export
 #' @rdname predict.tab_pfn
-augment.tab_pfn <- function(x, new_data, type = NULL, ...) {
+augment.tab_pfn <- function(
+  x,
+  new_data,
+  type = NULL,
+  quantile_levels = NULL,
+  ...
+) {
   new_data <- tibble::new_tibble(new_data)
-  res <- predict(x, new_data, type = type)
+  res <- predict(x, new_data, type = type, quantile_levels = quantile_levels)
   res <- cbind(res, new_data)
   tibble::new_tibble(res)
 }

tests/testthat/_snaps/classification.md

@@ -52,6 +52,10 @@
       Device:
       i cpu
 
+---
+
+    `quantile_levels` is only for regression models.
+
 # classification models - recipes
 
     Code

tests/testthat/_snaps/regression.md

@@ -53,6 +53,10 @@
 
     `tab_pfn()` is not defined for the number 1.
 
+# quantile regression models
+
+    `quantile_levels` must be a number between 0 and 1, not the number 1.9.
+
 # regression models - recipes
 
     Code

tests/testthat/test-classification.R

@@ -76,6 +76,8 @@ test_that('classification models', {
   expect_s3_class(aug_mat, c("tbl_df", "tbl", "data.frame"))
   expect_equal(nrow(aug_mat), 3L)
   expect_equal(ncol(aug_mat), 5L)
+
+  expect_snapshot_error(predict(mod_mat, x_te_mat, quantile_levels = 0.5))
 })
 
 test_that('classification models - recipes', {

tests/testthat/test-regression.R

@@ -74,6 +74,36 @@ test_that('regression models', {
   )
 })
 
+test_that('quantile regression models', {
+  skip_if_no_tabpfn()
+
+  quantile_levels <- c(0.1, 0.5, 0.9)
+  pred_ptype <- tibble::tibble(
+    .pred = numeric(0),
+    .pred_quantile = hardhat::quantile_pred(
+      matrix(numeric(0), ncol = length(quantile_levels)),
+      quantile_levels
+    )
+  )
+
+  set.seed(166)
+  mod <- tab_pfn(predictors, outcome)
+
+  pred <- predict(mod, mtcars[1:3, -1], quantile_levels = quantile_levels)
+  expect_equal(pred[0, ], pred_ptype)
+  expect_equal(nrow(pred), 3L)
+  expect_equal(pred$.pred, predict(mod, mtcars[1:3, -1])$.pred)
+  expect_true(all(apply(as.matrix(pred$.pred_quantile), 1, diff) >= 0))
+
+  expect_no_error(predict(mod, mtcars[1:3, -1], quantile_levels = 0.5))
+
+  expect_snapshot_error(predict(mod, mtcars[1:3, -1], quantile_levels = 1.9))
+
+  aug <- augment(mod, mtcars[1:3, -1], quantile_levels = quantile_levels)
+  expect_equal(aug[, names(pred)], pred)
+  expect_equal(ncol(aug), 12L)
+})
+
 test_that('training_set_limit with data frame and matrix interfaces', {
   skip_if_no_tabpfn()
   skip_if_not_installed("recipes")