Update deprecated API usage (@consumer, GenAxisArray)

cboulay · cboulay · commit c93ca9e77572 · 2026-01-30T14:13:11.000-05:00
diff --git a/src/ezmsg/learn/process/sgd.py b/src/ezmsg/learn/process/sgd.py
@@ -2,9 +2,12 @@
 
 import ezmsg.core as ez
 import numpy as np
-from ezmsg.baseproc import GenAxisArray
-from ezmsg.sigproc.sampler import SampleMessage
-from ezmsg.util.generator import consumer
+from ezmsg.baseproc import (
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+    SampleMessage,
+    processor_state,
+)
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 from sklearn.exceptions import NotFittedError
@@ -13,103 +16,6 @@
 from ..util import ClassifierMessage
 
 
-@consumer
-def sgd_decoder(
-    alpha: float = 1.5e-5,
-    eta0: float = 1e-7,  # Lower than what you'd use for offline training.
-    loss: str = "squared_hinge",
-    label_weights: dict[str, float] | None = None,
-    settings_path: str | None = None,
-) -> typing.Generator[AxisArray | SampleMessage, ClassifierMessage | None, None]:
-    """
-    Passive Aggressive Classifier
-    Online Passive-Aggressive Algorithms <http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf>
-    K. Crammer, O. Dekel, J. Keshat, S. Shalev-Shwartz, Y. Singer - JMLR (2006)
-
-    Args:
-        alpha: Maximum step size (regularization)
-        eta0: The initial learning rate for the 'adaptive’ schedules.
-        loss: The loss function to be used:
-            hinge: equivalent to PA-I in the reference paper.
-            squared_hinge: equivalent to PA-II in the reference paper.
-        label_weights: An optional dictionary of label names and their relative weight.
-            e.g., {'Go': 31.0, 'Stop': 0.5}
-            If this is None then settings_path must be provided and the pre-trained model
-        settings_path: Path to the stored sklearn model pkl file.
-
-    Returns:
-        Generator that accepts `SampleMessage` for incremental training (`partial_fit`) and yields None,
-        or `AxisArray` for inference (`predict`) and yields a `ClassifierMessage`.
-    """
-    # pre-init inputs and outputs
-    msg_out = ClassifierMessage(data=np.array([]), dims=[""])
-
-    # State variables:
-
-    if settings_path is not None:
-        import pickle
-
-        with open(settings_path, "rb") as f:
-            model = pickle.load(f)
-            if label_weights is not None:
-                model.class_weight = label_weights
-            # Overwrite eta0, probably with a value lower than what was used online.
-            model.eta0 = eta0
-    else:
-        model = SGDClassifier(
-            loss=loss,
-            alpha=alpha,
-            penalty="elasticnet",
-            learning_rate="adaptive",
-            eta0=eta0,
-            early_stopping=False,
-            class_weight=label_weights,
-        )
-
-    b_first_train = True
-    # TODO: template_out
-
-    while True:
-        msg_in: AxisArray | SampleMessage = yield msg_out
-
-        msg_out = None
-        if type(msg_in) is SampleMessage:
-            # SampleMessage used for training.
-            if not np.any(np.isnan(msg_in.sample.data)):
-                train_sample = msg_in.sample.data.reshape(1, -1)
-                if b_first_train:
-                    model.partial_fit(
-                        train_sample,
-                        [msg_in.trigger.value],
-                        classes=list(label_weights.keys()),
-                    )
-                    b_first_train = False
-                else:
-                    model.partial_fit(train_sample, [msg_in.trigger.value])
-        elif msg_in.data.size:
-            # AxisArray used for inference
-            if not np.any(np.isnan(msg_in.data)):
-                try:
-                    X = msg_in.data.reshape((msg_in.data.shape[0], -1))
-                    result = model._predict_proba_lr(X)
-                except NotFittedError:
-                    result = None
-                if result is not None:
-                    out_axes = {}
-                    if msg_in.dims[0] in msg_in.axes:
-                        out_axes[msg_in.dims[0]] = replace(
-                            msg_in.axes[msg_in.dims[0]],
-                            offset=msg_in.axes[msg_in.dims[0]].offset,
-                        )
-                    msg_out = ClassifierMessage(
-                        data=result,
-                        dims=msg_in.dims[:1] + ["labels"],
-                        axes=out_axes,
-                        labels=list(model.class_weight.keys()),
-                        key=msg_in.key,
-                    )
-
-
 class SGDDecoderSettings(ez.Settings):
     alpha: float = 1e-5
     eta0: float = 3e-4
@@ -118,14 +24,96 @@ class SGDDecoderSettings(ez.Settings):
     settings_path: str | None = None
 
 
-class SGDDecoder(GenAxisArray):
-    SETTINGS = SGDDecoderSettings
-    INPUT_SAMPLE = ez.InputStream(SampleMessage)
+@processor_state
+class SGDDecoderState:
+    model: typing.Any = None
+    b_first_train: bool = True
 
-    # Method to be implemented by subclasses to construct the specific generator
-    def construct_generator(self):
-        self.STATE.gen = sgd_decoder(**self.SETTINGS.__dict__)
 
-    @ez.subscriber(INPUT_SAMPLE)
-    async def on_sample(self, msg: SampleMessage) -> None:
-        _ = self.STATE.gen.send(msg)
+class SGDDecoderTransformer(BaseAdaptiveTransformer[SGDDecoderSettings, AxisArray, ClassifierMessage, SGDDecoderState]):
+    """
+    SGD-based online classifier.
+
+    Online Passive-Aggressive Algorithms
+    <http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf>
+    K. Crammer, O. Dekel, J. Keshat, S. Shalev-Shwartz, Y. Singer - JMLR (2006)
+    """
+
+    def _refreshed_model(self):
+        if self.settings.settings_path is not None:
+            import pickle
+
+            with open(self.settings.settings_path, "rb") as f:
+                model = pickle.load(f)
+                if self.settings.label_weights is not None:
+                    model.class_weight = self.settings.label_weights
+                model.eta0 = self.settings.eta0
+        else:
+            model = SGDClassifier(
+                loss=self.settings.loss,
+                alpha=self.settings.alpha,
+                penalty="elasticnet",
+                learning_rate="adaptive",
+                eta0=self.settings.eta0,
+                early_stopping=False,
+                class_weight=self.settings.label_weights,
+            )
+        return model
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.model = self._refreshed_model()
+        self._state.b_first_train = True
+
+    def _process(self, message: AxisArray) -> ClassifierMessage | None:
+        if self._state.model is None or not message.data.size:
+            return None
+        if np.any(np.isnan(message.data)):
+            return None
+        try:
+            X = message.data.reshape((message.data.shape[0], -1))
+            result = self._state.model._predict_proba_lr(X)
+        except NotFittedError:
+            return None
+        out_axes = {}
+        if message.dims[0] in message.axes:
+            out_axes[message.dims[0]] = replace(
+                message.axes[message.dims[0]],
+                offset=message.axes[message.dims[0]].offset,
+            )
+        return ClassifierMessage(
+            data=result,
+            dims=message.dims[:1] + ["labels"],
+            axes=out_axes,
+            labels=list(self._state.model.class_weight.keys()),
+            key=message.key,
+        )
+
+    def partial_fit(self, message: SampleMessage) -> None:
+        if self._state.model is None:
+            # Initialize model on first training sample
+            self._state.model = self._refreshed_model()
+            self._state.b_first_train = True
+
+        if np.any(np.isnan(message.sample.data)):
+            return
+        train_sample = message.sample.data.reshape(1, -1)
+        if self._state.b_first_train:
+            self._state.model.partial_fit(
+                train_sample,
+                [message.trigger.value],
+                classes=list(self.settings.label_weights.keys()),
+            )
+            self._state.b_first_train = False
+        else:
+            self._state.model.partial_fit(train_sample, [message.trigger.value])
+
+
+class SGDDecoder(
+    BaseAdaptiveTransformerUnit[
+        SGDDecoderSettings,
+        AxisArray,
+        ClassifierMessage,
+        SGDDecoderTransformer,
+    ]
+):
+    SETTINGS = SGDDecoderSettings
diff --git a/tests/unit/test_sgd.py b/tests/unit/test_sgd.py
@@ -2,7 +2,7 @@
 from ezmsg.sigproc.sampler import SampleMessage, SampleTriggerMessage
 from ezmsg.util.messages.axisarray import AxisArray
 
-from ezmsg.learn.process.sgd import sgd_decoder
+from ezmsg.learn.process.sgd import SGDDecoderSettings, SGDDecoderTransformer
 
 
 def test_sgd():
@@ -46,14 +46,16 @@ def test_sgd():
     """
     label_weights = {k: 1.0 for k in time_idx.keys()}
     # Sending an axis array before it has seen any training samples should yield None
-    gen = sgd_decoder(alpha=1e-3, loss="squared_hinge", label_weights=label_weights)
-    assert gen.send(windows[0]) is None
+    decoder = SGDDecoderTransformer(
+        settings=SGDDecoderSettings(alpha=1e-3, loss="squared_hinge", label_weights=label_weights)
+    )
+    assert decoder(windows[0]) is None
 
     # Now let's try training on all samples
     for sample in samples:
-        gen.send(sample)
+        decoder(sample)
     # Then doing inference on all multi-wins
-    probas = [gen.send(win) for win in windows]
+    probas = [decoder(win) for win in windows]
 
     # With this easy-to-classify data, accuracy should be 100%
     # when we fit all training before predicting any test.
@@ -64,12 +66,14 @@ def test_sgd():
     assert np.array_equal(class_ids, expected_ids)
 
     # Try again (new model) but alternate 1 train, 2 test.
-    gen = sgd_decoder(alpha=1e-3, loss="squared_hinge", label_weights=label_weights)
+    decoder = SGDDecoderTransformer(
+        settings=SGDDecoderSettings(alpha=1e-3, loss="squared_hinge", label_weights=label_weights)
+    )
     probas = []
     for samp_ix, samp in enumerate(samples):
-        gen.send(samp)
-        probas.append(gen.send(windows[samp_ix * 2]))
-        probas.append(gen.send(windows[samp_ix * 2 + 1]))
+        decoder(samp)
+        probas.append(decoder(windows[samp_ix * 2]))
+        probas.append(decoder(windows[samp_ix * 2 + 1]))
     class_ids = []
     for cm in probas:
         class_ids.extend(np.argmax(cm.data, axis=1).tolist())
