diff --git a/DESCRIPTION b/DESCRIPTION index b49daaf..6e1305e 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -21,7 +21,7 @@ Imports: cli, dplyr, generics, - hardhat, + hardhat (>= 1.4.1), jsonlite, purrr, reticulate (>= 1.41.0.1), diff --git a/R/TabPFN-predict.R b/R/TabPFN-predict.R index fe62ecf..3247167 100644 --- a/R/TabPFN-predict.R +++ b/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) } diff --git a/man/predict.tab_pfn.Rd b/man/predict.tab_pfn.Rd index 0e28f57..db42b85 100644 --- a/man/predict.tab_pfn.Rd +++ b/man/predict.tab_pfn.Rd @@ -5,9 +5,9 @@ \alias{augment.tab_pfn} \title{Predict using \code{TabPFN}} \usage{ -\method{predict}{tab_pfn}(object, new_data, type = NULL, ...) +\method{predict}{tab_pfn}(object, new_data, type = NULL, quantile_levels = NULL, ...) -\method{augment}{tab_pfn}(x, new_data, type = NULL, ...) +\method{augment}{tab_pfn}(x, new_data, type = NULL, quantile_levels = NULL, ...) } \arguments{ \item{object, x}{A \code{tab_pfn} object.} @@ -17,6 +17,10 @@ \item{type}{The type of prediction. For classification, can be \code{"class"} or \code{"prob"}. Defaults to \code{NULL} which gives all prediction types possible.} +\item{quantile_levels}{A numeric vector of probabilities, sorted in +increasing order, at which to predict the outcome distribution. Regression +only; defaults to \code{NULL} for no quantile predictions.} + \item{...}{Not used, but required for extensibility.} } \value{ @@ -28,6 +32,9 @@ For regression data, the prediction is in the column \code{.pred}. For classification, the class predictions are in \code{.pred_class} and the probability estimates are in columns with the pattern \verb{.pred_\{level\}} where \code{level} is the levels of the outcome factor vector. + +When \code{quantile_levels} is given, regression results also have a +\code{.pred_quantile} column of \code{\link[hardhat:quantile_pred]{hardhat::quantile_pred()}} values. } \description{ Predict using \code{TabPFN} @@ -44,6 +51,7 @@ if (is_tab_pfn_installed() & interactive()) { # Predict predict(mod, car_test) + predict(mod, car_test, quantile_levels = c(0.1, 0.5, 0.9)) augment(mod, car_test) } } diff --git a/tests/testthat/_snaps/classification.md b/tests/testthat/_snaps/classification.md index ef23d56..2830217 100644 --- a/tests/testthat/_snaps/classification.md +++ b/tests/testthat/_snaps/classification.md @@ -52,6 +52,10 @@ Device: i cpu +--- + + `quantile_levels` is only for regression models. + # classification models - recipes Code diff --git a/tests/testthat/_snaps/regression.md b/tests/testthat/_snaps/regression.md index 9356109..b6fc930 100644 --- a/tests/testthat/_snaps/regression.md +++ b/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 diff --git a/tests/testthat/test-classification.R b/tests/testthat/test-classification.R index f23521c..db1704d 100644 --- a/tests/testthat/test-classification.R +++ b/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', { diff --git a/tests/testthat/test-regression.R b/tests/testthat/test-regression.R index 40c6fb7..a670c8f 100644 --- a/tests/testthat/test-regression.R +++ b/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")