cursus.steps.scripts.model_calibration

Model Calibration Script for SageMaker Processing.

This script calibrates model prediction scores to accurate probabilities, which is essential for risk-based decision-making and threshold setting. It supports multiple calibration methods including GAM, Isotonic Regression, and Platt Scaling, with options for monotonicity constraints.

Supported Scenarios: - Binary single-task: One score field with binary labels - Multi-class single-task: Multiple score fields (one per class) with categorical labels - Multi-task binary: Multiple independent binary tasks, each with its own score and label fields

Environment Variables:

Single-Task Binary:

SCORE_FIELD: Name of score column (e.g., “prob_class_1”) LABEL_FIELD: Name of label column (e.g., “label”) IS_BINARY: “true”

Multi-Task Binary (e.g., LightGBMMT):

SCORE_FIELDS: Comma-separated score columns (e.g., “task1_prob,task2_prob,task3_prob”) TASK_LABEL_NAMES: Comma-separated label columns (e.g., “task1_true,task2_true,task3_true”)

Optional - will be inferred from score field names if not provided

IS_BINARY: “true” (required)

Multi-Class Single-Task:

SCORE_FIELD_PREFIX: Prefix for probability columns (e.g., “prob_class_”) LABEL_FIELD: Name of label column NUM_CLASSES: Number of classes MULTICLASS_CATEGORIES: JSON array of class names IS_BINARY: “false”

Note: Multi-class multi-task calibration is not currently supported.

check_call(cmd, *a, **k)

install_packages_from_public_pypi(packages)

Install packages from standard public PyPI.

Parameters:

packages (list) – List of package specifications (e.g., [“pandas==1.5.0”, “numpy”])

install_packages_from_secure_pypi(packages)

Install packages from secure CodeArtifact PyPI.

Parameters:

packages (list) – List of package specifications (e.g., [“pandas==1.5.0”, “numpy”])

install_packages(packages, use_secure=False)

Install packages from PyPI source based on configuration.

This is the main installation function that delegates to either public or secure PyPI based on the USE_SECURE_PYPI environment variable.

Parameters:

• packages (list) – List of package specifications (e.g., [“pandas==1.5.0”, “numpy”])

• use_secure (bool) – If True, use secure CodeArtifact PyPI; if False, use public PyPI. Defaults to USE_SECURE_PYPI environment variable.

Environment Variables:

USE_SECURE_PYPI: Set to “true” to use secure PyPI, “false” for public PyPI

Example

# Install from public PyPI (default) install_packages([“pandas==1.5.0”, “numpy”])

# Install from secure PyPI os.environ[“USE_SECURE_PYPI”] = “true” install_packages([“pandas==1.5.0”, “numpy”])

load_dataframe_with_format(file_path)

Load DataFrame and detect its format.

Parameters:

file_path – Path to the file

Returns:

Tuple of (DataFrame, format_string)

Return type:

Tuple[DataFrame, str]

save_dataframe_with_format(df, output_path, format_str)

Save DataFrame in specified format.

Parameters:

• df (DataFrame) – DataFrame to save

• output_path – Base output path (without extension)

• format_str (str) – Format to save in (‘csv’, ‘tsv’, or ‘parquet’)

Returns:

Path to saved file

class CalibrationConfig(input_data_path='/opt/ml/processing/input/eval_data', output_calibration_path='/opt/ml/processing/output/calibration', output_metrics_path='/opt/ml/processing/output/metrics', output_calibrated_data_path='/opt/ml/processing/output/calibrated_data', calibration_method='gam', label_field='label', score_field='prob_class_1', is_binary=True, monotonic_constraint=True, gam_splines=10, error_threshold=0.05, num_classes=2, score_field_prefix='prob_class_', multiclass_categories=None, calibration_sample_points=1000)

Bases: object

Configuration class for model calibration.

classmethod from_env()

Create configuration from environment variables.

create_directories(config)

Create output directories if they don’t exist.

find_first_data_file(data_dir=None, config=None)

Find the first supported data file in directory.

Parameters:

• data_dir (str | None) – Directory to search for data files (defaults to config input_data_path)

• config (CalibrationConfig) – Configuration object (required)

Returns:

Path to the first supported data file found

Return type:

str

Raises:

FileNotFoundError – If no supported data file is found

load_data(config, is_multitask=False)

Load evaluation data with predictions using format preservation.

Parameters:

• config (CalibrationConfig) – Configuration object (required)

• is_multitask (bool) – If True, skip score_field validation (multi-task mode)

Returns:

Loaded evaluation data and detected format

Return type:

Tuple[pd.DataFrame, str]

Raises:

• FileNotFoundError – If no data file is found

• ValueError – If required columns are missing

log_section(title)

Log a section title with delimiters for better visibility.

parse_score_fields(environ_vars)

Parse SCORE_FIELD or SCORE_FIELDS from environment variables.

Parameters:

environ_vars (dict) – Dictionary of environment variables

Returns:

List of score field names to calibrate

Raises:

ValueError – If neither SCORE_FIELD nor SCORE_FIELDS is provided

Return type:

List[str]

parse_task_label_fields(environ_vars, score_fields)

Parse TASK_LABEL_NAMES or infer from score_fields.

For multi-task calibration, each score field needs a corresponding label field. This function either: 1. Uses explicit TASK_LABEL_NAMES environment variable 2. Infers labels from score field names (_prob -> _true) 3. Falls back to single LABEL_FIELD for backward compatibility

Parameters:

• environ_vars (dict) – Dictionary of environment variables

• score_fields (List[str]) – List of score field names

Returns:

List of label field names, one per score field

Raises:

ValueError – If TASK_LABEL_NAMES length doesn’t match score_fields

Return type:

List[str]

validate_score_fields(df, score_fields, label_field)

Validate that score fields exist in the DataFrame.

Parameters:

• df (DataFrame) – Input DataFrame

• score_fields (List[str]) – List of score field names

• label_field (str) – Name of the label field

Returns:

List of valid score fields (fields that exist in DataFrame)

Raises:

ValueError – If no valid score fields are found

Return type:

List[str]

extract_and_load_nested_tarball_data(config)

Extract and load data from nested tar.gz files in SageMaker output structure.

Handles SageMaker’s specific output structure: - output.tar.gz (outer archive)

• val.tar.gz (inner archive) - val/predictions.csv (actual data) - val_metrics/… (metrics and plots)

• test.tar.gz (inner archive) - test/predictions.csv (actual data) - test_metrics/… (metrics and plots)

Also handles cases where the input path contains: - Direct output.tar.gz file - Path to a job directory that contains output/output.tar.gz - Path to a parent directory with job subdirectories

Parameters:

config (CalibrationConfig) – Configuration object (required)

Returns:

Combined dataset with predictions from extracted tar.gz files

Return type:

pd.DataFrame

Raises:

FileNotFoundError – If necessary tar.gz files or prediction data not found

load_and_prepare_data(config, job_type='calibration', is_multitask=False)

Load evaluation data and prepare it for calibration based on classification type.

Parameters:

• config (CalibrationConfig) – Configuration object (required)

• job_type (str) – The job type to determine how to load data

• is_multitask (bool) – If True, skip score_field validation (multi-task mode)

Returns:

Different return values based on classification type:

• Multi-task: (df, None, None, None) - main function handles per-task extraction

• Binary: (df, y_true, y_prob, None)

• Multi-class: (df, y_true, None, y_prob_matrix)

Return type:

tuple

Raises:

• FileNotFoundError – If no data file is found

• ValueError – If required columns are missing

train_gam_calibration(scores, labels, config)

Train a GAM calibration model and convert to lookup table.

Parameters:

• scores (ndarray) – Raw prediction scores to calibrate

• labels (ndarray) – Ground truth binary labels (0/1)

• config (CalibrationConfig) – Configuration object (required)

Returns:

Lookup table as [(raw_score, calibrated_score), …]

Same format as percentile calibration for unified inference code.

Return type:

List[Tuple[float, float]]

Raises:

ImportError – If pygam is not installed

train_isotonic_calibration(scores, labels, config)

Train an isotonic regression calibration model and convert to lookup table.

Parameters:

• scores (ndarray) – Raw prediction scores to calibrate

• labels (ndarray) – Ground truth binary labels (0/1)

• config (CalibrationConfig) – Configuration object (required)

Returns:

Lookup table as [(raw_score, calibrated_score), …]

Same format as percentile calibration for unified inference code.

Return type:

List[Tuple[float, float]]

train_platt_scaling(scores, labels, config)

Train a Platt scaling (logistic regression) calibration model and convert to lookup table.

Parameters:

• scores (ndarray) – Raw prediction scores to calibrate

• labels (ndarray) – Ground truth binary labels (0/1)

• config (CalibrationConfig) – Configuration object (required)

Returns:

Lookup table as [(raw_score, calibrated_score), …]

Same format as percentile calibration for unified inference code.

Return type:

List[Tuple[float, float]]

train_multiclass_calibration(y_prob_matrix, y_true, method='isotonic', config=None)

Train calibration models for each class in one-vs-rest fashion.

Parameters:

• y_prob_matrix (ndarray) – Matrix of prediction probabilities (samples × classes)

• y_true (ndarray) – Ground truth class labels

• method (str) – Calibration method to use (“gam”, “isotonic”, “platt”)

• config (CalibrationConfig) – Configuration object (required)

Returns:

List of calibration models, one for each class

Return type:

list

apply_multiclass_calibration(y_prob_matrix, calibrators, config)

Apply calibration to each class probability and normalize.

Parameters:

• y_prob_matrix (ndarray) – Matrix of uncalibrated probabilities (samples × classes)

• calibrators (List[Any]) – List of calibration models or lookup tables, one for each class

• config (CalibrationConfig) – Configuration object (required)

Returns:

Matrix of calibrated probabilities (samples × classes)

Return type:

np.ndarray

compute_calibration_metrics(y_true, y_prob, n_bins=10)

Compute comprehensive calibration metrics including ECE, MCE, and reliability diagram.

This function calculates: - Expected Calibration Error (ECE): weighted average of absolute calibration errors - Maximum Calibration Error (MCE): maximum calibration error across all bins - Reliability diagram data: points for plotting calibration curve - Bin statistics: detailed information about each probability bin - Brier score: quadratic scoring rule for probabilistic predictions - Preservation of discrimination: comparison of AUC before/after calibration

Parameters:

• y_true (ndarray) – Ground truth binary labels (0/1)

• y_prob (ndarray) – Predicted probabilities

• n_bins (int) – Number of bins for calibration curve

Returns:

Dictionary containing calibration metrics

Return type:

Dict

compute_multiclass_calibration_metrics(y_true, y_prob_matrix, n_bins=10, config=None)

Compute calibration metrics for multi-class scenario.

Parameters:

• y_true (ndarray) – Ground truth class labels

• y_prob_matrix (ndarray) – Matrix of prediction probabilities (samples × classes)

• n_bins (int) – Number of bins for calibration curve

• config (CalibrationConfig) – Configuration object (required)

Returns:

Dictionary containing calibration metrics

Return type:

dict

plot_reliability_diagram(y_true, y_prob_uncalibrated, y_prob_calibrated, n_bins=10, config=None)

Create reliability diagram comparing uncalibrated and calibrated probabilities.

Parameters:

• y_true (ndarray) – Ground truth binary labels (0/1)

• y_prob_uncalibrated (ndarray) – Uncalibrated prediction probabilities

• y_prob_calibrated (ndarray) – Calibrated prediction probabilities

• n_bins (int) – Number of bins for calibration curve

• config (CalibrationConfig) – Configuration object (required)

Returns:

Path to the saved figure

Return type:

str

plot_multiclass_reliability_diagram(y_true, y_prob_uncalibrated, y_prob_calibrated, n_bins=10, config=None)

Create reliability diagrams for multi-class case, one plot per class.

Parameters:

• y_true – Ground truth class labels

• y_prob_uncalibrated – Matrix of uncalibrated probabilities (samples × classes)

• y_prob_calibrated – Matrix of calibrated probabilities (samples × classes)

• n_bins – Number of bins for calibration curve

• config (CalibrationConfig) – Configuration object (required)

Returns:

Path to the saved figure

Return type:

str

main(input_paths, output_paths, environ_vars, job_args=None)

Main entry point for the calibration script.

Parameters:

• input_paths (dict) – Dictionary of input paths with logical names

• output_paths (dict) – Dictionary of output paths with logical names

• environ_vars (dict) – Dictionary of environment variables

• job_args (Namespace) – Command line arguments (optional)

Returns:

Dictionary with metrics and results

Return type:

dict