predict_proba

TabularFoundationModel.predict_proba(test_data: str | Path | DataFrame, train_data: str | Path | DataFrame, label: str, *, include_predict: bool = True, predictions_path: str | None = None, hyperparameters: Dict[str, Any] | None = None, instance_type: str | None = None, framework_version: str = 'latest', custom_image_uri: str | None = None, wait: bool = True, **backend_kwargs) → Tuple[Series, DataFrame | Series] | DataFrame | Series | JobPredictionFuture[source]

Run batch prediction returning class probabilities.

Identical to predict() but returns class probabilities. For regression the probabilities are identical to the predictions.

Parameters:

  • test_data – Data to predict on. Must contain every feature column present in train_data except label.

  • train_data – Labeled few-shot context for the foundation model, as a DataFrame or local/S3 path to a data file.

  • label – Target column name in train_data.

  • include_predict – Whether to return the predictions along with the probabilities. Comes for free — the job always computes both.

  • predictions_path – S3 URL where predictions will be written by the training container. Defaults to {cloud_output_path}/{job_name}/predictions.csv.

  • hyperparameters – Model hyperparameters for inference. Overrides values passed to the constructor.

  • instance_type – Instance type for the prediction job. If None, uses registry default.

  • framework_version – Container framework version.

  • custom_image_uri – Custom Docker image URI for the container.

  • wait – If True, block and return the result. If False, return a JobPredictionFuture immediately.

  • **backend_kwargs – Additional backend-specific arguments (e.g., job_name, volume_size).

Returns:

If include_predict is True, returns (prediction, predict_probability); otherwise just predict_probability. Returns a JobPredictionFuture when wait=False.

Return type:

(pd.Series, pd.DataFrame | pd.Series) or (pd.DataFrame | pd.Series) or JobPredictionFuture