Skip to content

Add predict_proba output option to SklearnModelProcessor#15

Merged
cboulay merged 2 commits into
ezmsg-org:devfrom
santiemckenzie:feat/sklearn-predict-proba
Jul 8, 2026
Merged

Add predict_proba output option to SklearnModelProcessor#15
cboulay merged 2 commits into
ezmsg-org:devfrom
santiemckenzie:feat/sklearn-predict-proba

Conversation

@santiemckenzie

Copy link
Copy Markdown

Add a predict_method setting ("predict" | "predict_proba", default "predict") so the processor can emit a per-class posterior vector instead of only a hard label. When set to "predict_proba", the output ch axis is relabeled to the model's classes_. Supports sklearn predict_proba and River predict_proba_many. The default predict path is unchanged.

Motivated by consumers that pool per-bin posteriors over a window (see BlackrockNeurotech/intent-pipelines#13), which previously required a bespoke decode unit because SklearnModelProcessor only exposed predict.

@cboulay

cboulay commented Jul 8, 2026

Copy link
Copy Markdown
Member

Thanks for the PR — this is well put together. The test coverage is appreciated, especially the default-behavior regression test, and the motivation makes sense.

One requested change before merge: please convert predict_method from a bare str to a (str, Enum), matching the existing convention in src/ezmsg/learn/util.py (RegressorType, AdaptiveLinearRegressor, etc.). This keeps the door open for future inference methods (decision_function, predict_log_proba, ...) without an API break, and because it subclasses str, plain-string configs keep working.

In process/sklearn.py:

from enum import Enum


class PredictMethod(str, Enum):
    PREDICT = "predict"
    PREDICT_PROBA = "predict_proba"

Settings:

predict_method: PredictMethod = PredictMethod.PREDICT

Comparisons then read as:

if self.settings.predict_method == PredictMethod.PREDICT_PROBA:

which also works when the setting was supplied as a plain string like "predict_proba". With that in place, the hand-rolled membership check in _reset_state can be reduced to a one-line coercion (PredictMethod(self.settings.predict_method)) or dropped.

One thing I noticed that's out of scope for this PR: with River classifiers the ch axis falls back to np.arange labels even though predict_proba_many returns the actual class labels as DataFrame columns. I've opened #17 to track that — no action needed here.

@santiemckenzie

Copy link
Copy Markdown
Author

Made those changes: predict_method is now a (str, Enum), and the check in _reset_state is down to the one-line coercion. Thank you!

@cboulay cboulay merged commit 7828f31 into ezmsg-org:dev Jul 8, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants