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 = '1.6', wait: bool = True, backend_overrides: dict[str, dict[str, Any]] | None = None, **kwargs: Unpack[TrainingJobKwargs]) → tuple[Series, DataFrame | Series] | DataFrame | Series | JobPredictionFuture[source]

Predict class probabilities for test_data in a one-off SageMaker job, using the labeled train_data as in-context examples.

For regression, the probability result is identical to the prediction.

Parameters:

  • test_data (str | Path | pd.DataFrame) – Rows to predict, as a pd.DataFrame or local/S3 path to a data file. Must contain all feature columns of train_data.

  • train_data (str | Path | pd.DataFrame) – Labeled examples, as a pd.DataFrame or local/S3 path to a data file.

  • label (str) – Name of the label column in train_data.

  • include_predict (bool, default = True) – Whether to return the predictions along with the probabilities. Both are computed in the same job.

  • predictions_path (str | None, default = None) – S3 URL ending in .csv or .parquet where the job writes the predictions, e.g. s3://my-bucket/predictions.csv. Defaults to {cloud_output_path}/{job_name}/predictions.csv.

  • hyperparameters (dict[str, Any] | None, default = None) – Model hyperparameters for inference. Overrides values passed to the constructor.

  • instance_type (str | None, default = None) – Instance type for the prediction job. Defaults to the model registry value.

  • framework_version (str, optional) – AutoGluon version, e.g. “1.6”. Uses the official AutoGluon DLC image for this version.

  • wait (bool, default = True) – If True, block until the job completes and return the result. If False, return a JobPredictionFuture immediately; call its .status() to check progress and .result() to get the result.

Returns:

tuple[pd.Series, pd.DataFrame | pd.Series] | pd.DataFrame | pd.Series | JobPredictionFuture – (prediction, predict_probability) if include_predict is True, otherwise predict_probability. A JobPredictionFuture if wait=False.

Other Parameters:

  • backend_overrides (dict[str, dict[str, Any]] | None, default = None) – Raw SageMaker request fields for settings without a dedicated argument.

    • Keys: request names from the SageMaker API section below.

    • Values: request fields in PascalCase, as in the SageMaker API and boto3. Deep-merged over the request built by AutoGluon-Cloud; lists and other non-dict values replace the generated ones.

    • Example: {"CreateTrainingJob": {"RetryStrategy": {"MaximumRetryAttempts": 2}}}

  • job_name (str | None, default = None) – Name of the training job that runs the prediction. Auto-generated if not set.

  • volume_size (int, default = 100) – Size in GB of the storage volume to use for the job.

  • custom_image_uri (str | None, default = None) – Custom container image URI. If set, framework_version is ignored.

  • timeout (int, default = 86400) – Maximum job runtime in seconds. Defaults to 24 hours.

SageMaker API

  • CreateTrainingJob: runs the prediction as a training job (not a batch transform job) on instance_type. Predictions are written to predictions_path.