Skip to main content
The TabPFN MCP server exposes a number of tools for performing classification and regression on Prior Labs’ managed GPU infrastructure:

upload_dataset

Get a secure upload URL for your dataset. We recommend this for most workflows: your data is sent directly to cloud storage instead of through the chat, so the agent can handle larger datasets without running into context limits or long execution time.
Uploading the file to the URL requires a sandbox execution environment and outbound network access. Not all MCP clients support this, and some may require a paid plan.
This tool returns a dataset_id and upload_url - valid for 60 minutes. Call this tool separately for uploading the training set and test dataset.

Required Parameters

filename
String
required
The filename for the dataset. Must be train.csv for training data or test.csv for test data.
  • "train.csv" — Training data
  • "test.csv" — Test data, predictions will be generated for

fit_and_predict_from_dataset

Fit the TabPFN-2.5 model on your pre-uploaded data and generate predictions. Use this tool when you want to fit a new model from scratch. Upload both dataset CSV files with upload_dataset first, then pass the two dataset IDs here.

Required Parameters

train_dataset_id
String
required
The dataset ID you got from upload_dataset for your training CSV. That file should include all input columns plus the column you want to predict (the target).
test_dataset_id
String
required
The dataset ID you got from upload_dataset for your test CSV. It should have the same input columns as the training file, but not the target column.
target_column
String
required
The name of the target column in the training dataset — e.g. "price" or "churned".
task_type
Literal
required
The type of the predictive task.
  • "classification" — Predict a category or class, including probability distributions
  • "regression" — Predict a continuous value

Optional Parameters

output_type
String
default:"null"
Prediction output type, default "preds" for classification and "mean" for regression.

Returns

model_id
String
Unique ID for the fitted model. Save this value to reuse the model later with predict tools.
predictions
Array
Prediction results in the format specified by output_type. For classification with "preds", returns a 1D array of class labels. With "probas", returns a 2D array of class probabilities. For regression with "mean", returns a 1D array of predicted values.

predict_from_dataset

Run predictions with a previously fitted model on a new test set. Use upload_dataset to upload your new test dataset file, then pass that dataset_id and the model_id you got from a previous fit_and_predict_* call. The test CSV must have the same set of features the model was fitted on.

Required Parameters

model_id
String
required
The model ID of the previously fitted model, from fit_and_predict_from_dataset or fit_and_predict_inline.
test_dataset_id
String
required
The dataset ID of the test dataset predictions will be made from.
task_type
Literal
required
The type of the predictive task.
  • "classification" — Predict a category or class, including probability distributions
  • "regression" — Predict a continuous value

Optional Parameters

output_type
String
default:"null"
Prediction output type, default "preds" for classification and "mean" for regression.

Returns

model_id
String
Echo of the model ID used for prediction.
predictions
Array
Prediction results in the format specified by output_type. For classification with "preds", returns a 1D array of class labels. With "probas", returns a 2D array of class probabilities. For regression with "mean", returns a 1D array of predicted values.

fit_and_predict_inline

Use this tool when you want to fit a new model from scratch. It fits on your data and immediately returns predictions for your test set, along with a model_id for future reuse.
Best for small datasets that fit in the conversation without running into context limits.

Required Parameters

X_train
Matrix
required
Training features as a 2D array where rows represent samples and columns represent features.
  • Shape: (n_train_samples, n_features)
  • Data types: Numeric (int/float) or categorical (string) values
  • Flexibility: Handles missing values, outliers, and mixed data types automatically
X_train = [
  [1.5, "red", 3],
  [2.0, "blue", 4],
  [1.8, "red", 5]
]
y_train
Vector
required
Training targets as a 1D array that must align with X_train rows.
  • Shape: (n_train_samples,)
  • Classification: Class labels (e.g., [0, 1, 0] or ["cat", "dog", "cat"])
  • Regression: Numeric values (e.g., [23.5, 45.2, 12.8])
X_test
Matrix
required
Test features as a 2D array for generating predictions.
  • Shape: (n_test_samples, n_features)
  • Critical: Must have the same number of features as X_train
X_test = [
  [1.8, "red", 5],
  [2.3, "green", 2]
]
task_type
Literal
required
The type of the predictive task.
  • "classification" — Predict a category or class, including probability distributions
  • "regression" — Predict a continuous value

Optional Parameters

output_type
String
default:"null"
Prediction output type, default "preds" for classification and "mean" for regression.

Returns

model_id
String
Unique ID for the fitted model. Save this value to reuse the model later with predict tools.
predictions
Array
Prediction results in the format specified by output_type. For classification with "preds", returns a 1D array of class labels. With "probas", returns a 2D array of class probabilities. For regression with "mean", returns a 1D array of predicted values.

predict

Generate new predictions using a previously fitted TabPFN model. Use this tool after calling fit_and_predict_* to make predictions on new data using an existing model.
Best for small datasets that fit in the conversation without running into context limits.

Required Parameters

model_id
String
required
ID of a previously fitted model, returned from fit_and_predict_*.
model_id = "9f1526b2-388b-4849-b965-6373d35f1a6b"
X_test
Matrix
required
Test features as a 2D array for generating predictions.
  • Shape: (n_test_samples, n_features)
  • Critical: Must have the same number of features the model was originally fitted on
X_test = [
  [1.2, "red", 7],
  [3.4, "blue", 1]
]
task_type
Literal
required
The type of the predictive task.
  • "classification" — Predict a category or class, including probability distributions
  • "regression" — Predict a continuous value

Optional Parameters

output_type
String
default:"null"
Prediction output type, default "preds" for classification and "mean" for regression.

Returns

model_id
String
Echo of the model ID used for prediction.
predictions
Array
Prediction results in the format specified by output_type. Same structure as fit_and_predict_inline.