feat: add probability-aware evaluation metrics and reporting#178
feat: add probability-aware evaluation metrics and reporting#178marcellodebernardi wants to merge 2 commits into
Conversation
Greptile SummaryThis PR enhances model evaluation by adding probability-aware metrics and reporting capabilities across the evaluation pipeline. Key Changes
Testing
Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| plexe/helpers.py | Added compute_metric_proba function for probability-based metrics (ROC AUC, log loss, Brier score) and updated evaluate_on_sample to use probabilities when available |
| plexe/agents/model_evaluator.py | Added Phase 1.5 (Probability Analysis) to compute ROC curves, calibration curves, and Brier scores with matplotlib visualization support |
| plexe/models.py | Extended CoreMetricsReport dataclass with probability-aware fields: roc_auc_from_proba, brier_score, calibration_data, and roc_curve_data |
| plexe/templates/inference/xgboost_predictor.py | Added predict_proba method that returns DataFrame with labeled probability columns (e.g., proba_yes, proba_no) |
| plexe/templates/inference/pytorch_predictor.py | Added predict_proba method with sigmoid for binary and softmax for multiclass, handles sparse matrices correctly |
| tests/unit/test_proba_metrics.py | Comprehensive tests for compute_metric_proba covering binary/multiclass ROC AUC, log loss, Brier score, and DataFrame column ordering |
Last reviewed commit: 8edb5ab
There was a problem hiding this comment.
Pull request overview
This PR adds comprehensive probability-aware evaluation capabilities to the ML model evaluation workflow. The implementation introduces predict_proba methods across all predictor templates, a new compute_metric_proba function for probability-based metrics (ROC AUC, log loss, Brier score), automatic detection and use of probability predictions in evaluate_on_sample, and a new "Probability Analysis" phase in the model evaluator agent that generates calibration curves, ROC curves, and related diagnostic plots.
Changes:
- Added
predict_probamethods to all predictor templates (XGBoost, LightGBM, CatBoost, Keras, PyTorch) with consistent column naming and probability normalization - Implemented
compute_metric_probafunction with support for ROC AUC (binary/multiclass OVR/OVO), log loss, and Brier score with DataFrame-aware label extraction - Extended
CoreMetricsReportdataclass with optional probability-specific fields (roc_auc_from_proba, brier_score, calibration_data, roc_curve_data) - Integrated probability-aware evaluation into model evaluator workflow with a new Phase 1.5 for calibration and discrimination analysis
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_proba_metrics.py | New test suite for compute_metric_proba covering binary/multiclass scenarios and DataFrame column ordering |
| tests/unit/test_predictors_proba.py | New test suite verifying predict_proba implementations across all predictor types with dummy models |
| tests/unit/test_evaluate_on_sample.py | New tests verifying probability-based metric prioritization and graceful fallback behavior |
| tests/CODE_INDEX.md | Updated documentation index with new test files |
| plexe/tools/submission.py | Extended register_core_metrics_report with optional probability-related parameters |
| plexe/templates/inference/xgboost_predictor.py | Added predict_proba method with label encoder support and binary probability expansion |
| plexe/templates/inference/pytorch_predictor.py | Added predict_proba with sigmoid/softmax activation for raw logits |
| plexe/templates/inference/lightgbm_predictor.py | Added predict_proba method with label encoder support and binary probability expansion |
| plexe/templates/inference/keras_predictor.py | Added predict_proba extracting raw model outputs with binary probability expansion |
| plexe/templates/inference/catboost_predictor.py | Added predict_proba method with label encoder support and binary probability expansion |
| plexe/models.py | Added four optional probability-related fields to CoreMetricsReport dataclass |
| plexe/helpers.py | Added compute_metric_proba function and enhanced evaluate_on_sample to prioritize probability-based computation |
| plexe/agents/model_evaluator.py | Added Phase 1.5 Probability Analysis with matplotlib support and LLM prompts for calibration/ROC curve generation |
| plexe/CODE_INDEX.md | Updated documentation index with new helper functions and predictor methods |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def test_tree_predictors_predict_proba_binary(predictor_cls): | ||
| proba = np.array([[0.2, 0.8], [0.6, 0.4]]) | ||
| predictor = predictor_cls.__new__(predictor_cls) | ||
| predictor.model = _DummyProbaModel(proba) | ||
| predictor.pipeline = _DummyPipeline() | ||
| predictor.label_encoder = _DummyEncoder(["no", "yes"]) | ||
|
|
||
| df = predictor.predict_proba(pd.DataFrame({"f1": [1, 2], "f2": [3, 4]})) | ||
|
|
||
| assert list(df.columns) == ["proba_no", "proba_yes"] | ||
| _assert_proba_df(df, expected_rows=2, expected_cols=2) |
There was a problem hiding this comment.
The test functions in this file lack docstrings. Following the codebase convention (seen in test_helpers.py), test functions should have brief docstrings describing what is being tested. For example, test_tree_predictors_predict_proba_binary could have a docstring like "Test predict_proba for tree-based predictors with binary classification."
| def test_evaluate_on_sample_uses_proba(monkeypatch): | ||
| sample_df = pd.DataFrame({"feature": [1, 2, 3], "target": [0, 1, 0]}) | ||
| spark = _DummySpark(sample_df) | ||
|
|
||
| class DummyPredictor: | ||
| def __init__(self, model_dir: str): | ||
| pass | ||
|
|
||
| def predict(self, x): | ||
| return pd.DataFrame({"prediction": [0, 0, 0]}) | ||
|
|
||
| def predict_proba(self, x): | ||
| return pd.DataFrame({"proba_0": [0.8, 0.2, 0.7], "proba_1": [0.2, 0.8, 0.3]}) | ||
|
|
||
| import plexe.templates.inference.xgboost_predictor as xgb_module | ||
|
|
||
| monkeypatch.setattr(xgb_module, "XGBoostPredictor", DummyPredictor) | ||
| monkeypatch.setattr(helpers, "compute_metric_proba", lambda *args, **kwargs: 0.77) | ||
| monkeypatch.setattr(helpers, "compute_metric", lambda *args, **kwargs: 0.11) | ||
|
|
||
| result = evaluate_on_sample( | ||
| spark=spark, | ||
| sample_uri="dummy", | ||
| model_artifacts_path=Path("."), | ||
| model_type=ModelType.XGBOOST, | ||
| metric="roc_auc", | ||
| target_columns=["target"], | ||
| ) | ||
|
|
||
| assert result == 0.77 |
There was a problem hiding this comment.
The test functions in this file lack docstrings. Following the codebase convention (seen in test_helpers.py), test functions should have brief docstrings describing what is being tested. For example, test_evaluate_on_sample_uses_proba could have a docstring like "Test that evaluate_on_sample uses predict_proba for probability-based metrics."
| def test_compute_metric_proba_brier_score_multiclass_dataframe_order(): | ||
| y_true = np.array(["b", "a", "c", "b"]) | ||
| y_proba = np.array( | ||
| [ | ||
| [0.7, 0.2, 0.1], | ||
| [0.1, 0.8, 0.1], | ||
| [0.2, 0.3, 0.5], | ||
| [0.6, 0.2, 0.2], | ||
| ] | ||
| ) | ||
| proba_df = pd.DataFrame(y_proba, columns=["proba_b", "proba_a", "proba_c"]) | ||
|
|
||
| expected_one_hot = np.array( | ||
| [ | ||
| [1.0, 0.0, 0.0], | ||
| [0.0, 1.0, 0.0], | ||
| [0.0, 0.0, 1.0], | ||
| [1.0, 0.0, 0.0], | ||
| ] | ||
| ) | ||
| expected = np.mean(np.sum((y_proba - expected_one_hot) ** 2, axis=1)) | ||
|
|
||
| result = compute_metric_proba(y_true, proba_df, "brier_score") | ||
|
|
||
| assert result == pytest.approx(expected) |
There was a problem hiding this comment.
The test creates a DataFrame with columns ordered as ["proba_b", "proba_a", "proba_c"], but the expected_one_hot is constructed with the assumption that classes are ordered as ["b", "a", "c"]. However, the compute_metric_proba function extracts labels from column names and then creates a class_to_index mapping based on these extracted labels. The order of classes used for one-hot encoding should match the order of the probability columns in the input DataFrame. This test may be verifying the correct behavior, but it would be clearer if it included a comment explaining that it's testing the column-order-aware behavior of compute_metric_proba.
| except (AttributeError, ValueError, TypeError) as exc: | ||
| logger.warning( | ||
| "Probability metric computation failed (%s); falling back to labels. Error: %s", | ||
| metric, | ||
| exc, | ||
| exc_info=True, | ||
| ) |
There was a problem hiding this comment.
The exception handling catches AttributeError, ValueError, and TypeError broadly with exc_info=True. While this provides good diagnostics, it might mask bugs in the compute_metric_proba implementation by silently falling back to label-based metrics. Consider whether certain exception types (like AttributeError from a missing method) should fail fast rather than fall back, or add more specific error messages to distinguish between different failure modes.
| def predict_proba(self, x: pd.DataFrame) -> pd.DataFrame: | ||
| """ | ||
| Predict per-class probabilities on input DataFrame. | ||
|
|
||
| Applies sigmoid for single-logit binary models, otherwise softmax. | ||
| """ | ||
| # Transform features through pipeline | ||
| x_transformed = self.pipeline.transform(x) | ||
|
|
||
| # Handle sparse matrix output (e.g. from OneHotEncoder, CountVectorizer) | ||
| if scipy.sparse.issparse(x_transformed): | ||
| x_transformed = x_transformed.toarray() | ||
|
|
||
| x_tensor = torch.tensor(np.array(x_transformed, dtype=np.float32)) | ||
|
|
||
| with torch.no_grad(): | ||
| raw_output = self.model(x_tensor) | ||
|
|
||
| raw_output = raw_output.detach().cpu() | ||
| if raw_output.ndim == 1: | ||
| raw_output = raw_output.unsqueeze(1) | ||
|
|
||
| if raw_output.shape[1] == 1: | ||
| proba_pos = torch.sigmoid(raw_output).squeeze(1) | ||
| probabilities = torch.stack([1 - proba_pos, proba_pos], dim=1) | ||
| else: | ||
| probabilities = torch.softmax(raw_output, dim=1) | ||
|
|
||
| probabilities = probabilities.numpy() | ||
| columns = [f"proba_{i}" for i in range(probabilities.shape[1])] | ||
| return pd.DataFrame(probabilities, columns=columns) |
There was a problem hiding this comment.
The predict_proba method assumes that single-output PyTorch models output logits that need sigmoid activation for binary classification. However, if a model already applies sigmoid/softmax internally (e.g., in its forward method), this would result in double activation (sigmoid of sigmoid or softmax of softmax). Consider documenting this assumption or adding a check/warning if the raw_output values appear to already be probabilities (e.g., all values in [0, 1] range for single output).
| def test_compute_metric_proba_roc_auc_binary(): | ||
| y_true = np.array([0, 0, 1, 1]) | ||
| y_proba = np.array([0.1, 0.4, 0.35, 0.8]) | ||
|
|
||
| expected = roc_auc_score(y_true, y_proba) | ||
| result = compute_metric_proba(y_true, y_proba, "roc_auc") | ||
|
|
||
| assert result == pytest.approx(expected) |
There was a problem hiding this comment.
The test functions in this file lack docstrings. Following the codebase convention (seen in test_helpers.py), test functions should have brief docstrings describing what is being tested. For example, test_compute_metric_proba_roc_auc_binary could have a docstring like "Test ROC AUC computation with binary classification probabilities."
| labels = None | ||
| if isinstance(y_proba, pd.DataFrame): | ||
| proba_df = y_proba | ||
| proba = proba_df.values | ||
| if all(col.startswith("proba_") for col in proba_df.columns): | ||
| label_strs = [col[len("proba_") :] for col in proba_df.columns] | ||
| if np.issubdtype(y_true_arr.dtype, np.integer): | ||
| try: | ||
| labels = [int(label) for label in label_strs] | ||
| except ValueError: | ||
| labels = label_strs | ||
| elif np.issubdtype(y_true_arr.dtype, np.floating): | ||
| try: | ||
| labels = [float(label) for label in label_strs] | ||
| except ValueError: | ||
| labels = label_strs | ||
| else: | ||
| labels = label_strs |
There was a problem hiding this comment.
The labels extraction logic assumes probability columns follow the order in the DataFrame columns, but it doesn't verify this ordering matches the actual model output order. If the predictor's label_encoder.classes_ order doesn't match the probability column order in the DataFrame (which follows label_encoder.classes_), there could be a mismatch. Consider adding a validation check or documenting the assumption that the probability columns are always ordered consistently with model output.
| registry_key="_core_metrics_report", | ||
| ) | ||
| if not success: | ||
| logger.warning("Probability analysis phase failed - continuing with partial evaluation") |
There was a problem hiding this comment.
Phase 1.5 (Probability Analysis) reuses the same registry_key "_core_metrics_report" as Phase 1 (Core Metrics), which means it overwrites the initial core metrics report. While this appears intentional (the prompt instructs the agent to merge fields from the original report), there's a risk: if Phase 1.5 fails to properly copy all fields from core_metrics_report, data could be lost. Consider either using a different key and merging reports in the run method, or adding validation to ensure all required fields are preserved after Phase 1.5 completes.
| registry_key="_core_metrics_report", | |
| ) | |
| if not success: | |
| logger.warning("Probability analysis phase failed - continuing with partial evaluation") | |
| registry_key="_core_metrics_probability_report", | |
| ) | |
| if not success: | |
| logger.warning("Probability analysis phase failed - continuing with partial evaluation") | |
| else: | |
| # Merge Phase 1.5 results back into the core metrics report while preserving original fields | |
| proba_report = self.context.scratch.get("_core_metrics_probability_report") | |
| if isinstance(core_metrics, dict) and isinstance(proba_report, dict): | |
| # Preserve all original core metrics fields; allow Phase 1.5 to augment/override where needed | |
| merged_core_metrics = {**core_metrics, **proba_report} | |
| self.context.scratch["_core_metrics_report"] = merged_core_metrics | |
| elif proba_report is not None: | |
| # Fall back to using the probability report as-is if it's not a dict but still present | |
| self.context.scratch["_core_metrics_report"] = proba_report | |
| # If proba_report is None, keep the original core_metrics in "_core_metrics_report" |
| "IMPORTANT: In the probability analysis phase, if the predictor has predict_proba(), compute ROC " | ||
| "curve data, calibration curve, and Brier score. Use matplotlib with Agg backend to generate plots, " | ||
| "encode as base64 PNG, and include in the visualizations field.\\n" |
There was a problem hiding this comment.
The instructions about probability analysis are included in the general agent instructions for all phases, not just the Probability Analysis phase. This could be confusing for agents in other phases (Core Metrics, Error Analysis, etc.) who see instructions about predict_proba and matplotlib but aren't expected to use them. Consider moving these specific instructions to only appear in the Probability Analysis phase prompt instead of the general agent instructions.
|
Closing as this PR has some major issues. Will reopen from a fresh branch. |
This PR adds probability-aware model evaluation and reporting in the evaluator workflow. The aim is to improve classification evaluation fidelity by using predicted probabilities for probability-based metrics and by capturing calibration/discrimination artifacts in the core metrics report. It also adds predictor-level
predict_probasupport across inference templates and test coverage for the new behavior.Testing
poetry run pytest tests/unit/test_helpers.py tests/unit/test_proba_metrics.py tests/unit/test_evaluate_on_sample.py tests/unit/test_predictors_proba.pypoetry run ruff check plexe/helpers.py plexe/models.py plexe/tools/submission.py plexe/agents/model_evaluator.py plexe/templates/inference/xgboost_predictor.py plexe/templates/inference/lightgbm_predictor.py plexe/templates/inference/catboost_predictor.py plexe/templates/inference/keras_predictor.py plexe/templates/inference/pytorch_predictor.py tests/unit/test_evaluate_on_sample.py tests/unit/test_proba_metrics.py tests/unit/test_predictors_proba.py