ezmsg-learn 1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

ezmsg/learn/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0'
-__version_tuple__ = version_tuple = (1, 0)
+__version__ = version = '1.2.0'
+__version_tuple__ = version_tuple = (1, 2, 0)
 
 __commit_id__ = commit_id = None

ezmsg/learn/dim_reduce/adaptive_decomp.py

@@ -1,14 +1,14 @@
 import typing
 
-from sklearn.decomposition import IncrementalPCA, MiniBatchNMF
-import numpy as np
 import ezmsg.core as ez
-from ezmsg.sigproc.base import (
-    processor_state,
+import numpy as np
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
+    processor_state,
 )
 from ezmsg.util.messages.axisarray import AxisArray, replace
+from sklearn.decomposition import IncrementalPCA, MiniBatchNMF
 
 
 class AdaptiveDecompSettings(ez.Settings):
@@ -23,15 +23,11 @@ class AdaptiveDecompState:
     estimator: typing.Any = None
 
 
-EstimatorType = typing.TypeVar(
-    "EstimatorType", bound=typing.Union[IncrementalPCA, MiniBatchNMF]
-)
+EstimatorType = typing.TypeVar("EstimatorType", bound=typing.Union[IncrementalPCA, MiniBatchNMF])
 
 
 class AdaptiveDecompTransformer(
-    BaseAdaptiveTransformer[
-        AdaptiveDecompSettings, AxisArray, AxisArray, AdaptiveDecompState
-    ],
+    BaseAdaptiveTransformer[AdaptiveDecompSettings, AxisArray, AxisArray, AdaptiveDecompState],
     typing.Generic[EstimatorType],
 ):
     """
@@ -80,9 +76,7 @@ class AdaptiveDecompTransformer(
             it_ax_ix = message.get_axis_idx(iter_axis)
             # Remaining axes are to be treated independently
             off_targ_axes = [
-                _
-                for _ in (message.dims[:it_ax_ix] + message.dims[it_ax_ix + 1 :])
-                if _ != self.settings.axis
+                _ for _ in (message.dims[:it_ax_ix] + message.dims[it_ax_ix + 1 :]) if _ != self.settings.axis
             ]
         self._state.axis_groups = iter_axis, targ_axes, off_targ_axes
 
@@ -152,9 +146,7 @@ class AdaptiveDecompTransformer(
 
         # Transform data
         if hasattr(self._state.estimator, "components_"):
-            decomp_dat = self._state.estimator.transform(in_dat).reshape(
-                (-1,) + self._state.template.data.shape[1:]
-            )
+            decomp_dat = self._state.estimator.transform(in_dat).reshape((-1,) + self._state.template.data.shape[1:])
             replace_kwargs["data"] = decomp_dat
 
         return replace(self._state.template, **replace_kwargs)
@@ -241,9 +233,7 @@ class MiniBatchNMFTransformer(AdaptiveDecompTransformer[MiniBatchNMF]):
     pass
 
 
-SettingsType = typing.TypeVar(
-    "SettingsType", bound=typing.Union[IncrementalPCASettings, MiniBatchNMFSettings]
-)
+SettingsType = typing.TypeVar("SettingsType", bound=typing.Union[IncrementalPCASettings, MiniBatchNMFSettings])
 TransformerType = typing.TypeVar(
     "TransformerType",
     bound=typing.Union[IncrementalPCATransformer, MiniBatchNMFTransformer],

ezmsg/learn/dim_reduce/incremental_decomp.py

@@ -1,14 +1,14 @@
 import typing
 
-import numpy as np
 import ezmsg.core as ez
-from ezmsg.util.messages.axisarray import AxisArray, replace
-from ezmsg.sigproc.base import (
-    CompositeProcessor,
+import numpy as np
+from ezmsg.baseproc import (
     BaseStatefulProcessor,
     BaseTransformerUnit,
+    CompositeProcessor,
 )
 from ezmsg.sigproc.window import WindowTransformer
+from ezmsg.util.messages.axisarray import AxisArray, replace
 
 from .adaptive_decomp import (
     IncrementalPCASettings,
@@ -36,9 +36,7 @@ class IncrementalDecompSettings(ez.Settings):
     forget_factor: float = 0.7
 
 
-class IncrementalDecompTransformer(
-    CompositeProcessor[IncrementalDecompSettings, AxisArray, AxisArray]
-):
+class IncrementalDecompTransformer(CompositeProcessor[IncrementalDecompSettings, AxisArray, AxisArray]):
     """
     Automates usage of IncrementalPCATransformer and MiniBatchNMFTransformer by using a WindowTransformer
     to extract training samples then calls partial_fit on the decomposition transformer.
@@ -125,15 +123,11 @@ class IncrementalDecompTransformer(
             # If the estimator has not been trained once, train it with the first message
             self._procs["decomp"].partial_fit(message)
         elif "windowing" in self._procs:
-            state["windowing"], train_msg = self._procs["windowing"].stateful_op(
-                state.get("windowing", None), message
-            )
+            state["windowing"], train_msg = self._procs["windowing"].stateful_op(state.get("windowing", None), message)
             self._partial_fit_windowed(train_msg)
 
         # Process the incoming message
-        state["decomp"], result = self._procs["decomp"].stateful_op(
-            state.get("decomp", None), message
-        )
+        state["decomp"], result = self._procs["decomp"].stateful_op(state.get("decomp", None), message)
 
         return state, result
 
@@ -174,8 +168,6 @@ class IncrementalDecompTransformer(
 
 
 class IncrementalDecompUnit(
-    BaseTransformerUnit[
-        IncrementalDecompSettings, AxisArray, AxisArray, IncrementalDecompTransformer
-    ]
+    BaseTransformerUnit[IncrementalDecompSettings, AxisArray, AxisArray, IncrementalDecompTransformer]
 ):
     SETTINGS = IncrementalDecompSettings

ezmsg/learn/linear_model/adaptive_linear_regressor.py

@@ -1,6 +1,12 @@
 from ..process.adaptive_linear_regressor import (
     AdaptiveLinearRegressorSettings as AdaptiveLinearRegressorSettings,
+)
+from ..process.adaptive_linear_regressor import (
     AdaptiveLinearRegressorState as AdaptiveLinearRegressorState,
+)
+from ..process.adaptive_linear_regressor import (
     AdaptiveLinearRegressorTransformer as AdaptiveLinearRegressorTransformer,
+)
+from ..process.adaptive_linear_regressor import (
     AdaptiveLinearRegressorUnit as AdaptiveLinearRegressorUnit,
 )

ezmsg/learn/linear_model/linear_regressor.py

@@ -1,5 +1,9 @@
 from ..process.linear_regressor import (
     LinearRegressorSettings as LinearRegressorSettings,
+)
+from ..process.linear_regressor import (
     LinearRegressorState as LinearRegressorState,
+)
+from ..process.linear_regressor import (
     LinearRegressorTransformer as LinearRegressorTransformer,
 )

ezmsg/learn/linear_model/sgd.py

@@ -1,5 +1,9 @@
 from ..process.sgd import (
-    sgd_decoder as sgd_decoder,
-    SGDDecoderSettings as SGDDecoderSettings,
     SGDDecoder as SGDDecoder,
 )
+from ..process.sgd import (
+    SGDDecoderSettings as SGDDecoderSettings,
+)
+from ..process.sgd import (
+    sgd_decoder as sgd_decoder,
+)

ezmsg/learn/linear_model/slda.py

@@ -1,6 +1,12 @@
+from ..process.slda import (
+    SLDA as SLDA,
+)
 from ..process.slda import (
     SLDASettings as SLDASettings,
+)
+from ..process.slda import (
     SLDAState as SLDAState,
+)
+from ..process.slda import (
     SLDATransformer as SLDATransformer,
-    SLDA as SLDA,
 )

ezmsg/learn/model/mlp.py

@@ -27,11 +27,12 @@ class MLP(torch.nn.Module):
         Initialize the MLP model.
         Args:
             input_size (int): The size of the input features.
-            hidden_size (int | list[int]): The sizes of the hidden layers. If a list, num_layers must be None or the length
-                of the list. If a single integer, num_layers must be specified and determines the number of hidden layers.
+            hidden_size (int | list[int]): The sizes of the hidden layers. If a list, num_layers must be None or the
+                length of the list. If a single integer, num_layers must be specified and determines the number of
+                hidden layers.
             num_layers (int, optional): The number of hidden layers. Length of hidden_size if None. Default is None.
-            output_heads (int | dict[str, int], optional): Number of output features or classes if single head output or a
-                dictionary mapping head names to output sizes if multi-head output. Default is 2 (single head).
+            output_heads (int | dict[str, int], optional): Number of output features or classes if single head output
+                or a dictionary mapping head names to output sizes if multi-head output. Default is 2 (single head).
             norm_layer (str, optional): A normalization layer to be applied after each linear layer. Default is None.
                 Common choices are "BatchNorm1d" or "LayerNorm".
             activation_layer (str, optional): An activation function to be applied after each normalization
@@ -43,18 +44,14 @@ class MLP(torch.nn.Module):
         super().__init__()
         if isinstance(hidden_size, int):
             if num_layers is None:
-                raise ValueError(
-                    "If hidden_size is an integer, num_layers must be specified."
-                )
+                raise ValueError("If hidden_size is an integer, num_layers must be specified.")
             hidden_size = [hidden_size] * num_layers
         if len(hidden_size) == 0:
             raise ValueError("hidden_size must have at least one element")
         if any(not isinstance(x, int) for x in hidden_size):
             raise ValueError("hidden_size must contain only integers")
         if num_layers is not None and len(hidden_size) != num_layers:
-            raise ValueError(
-                "Length of hidden_size must match num_layers if num_layers is specified."
-            )
+            raise ValueError("Length of hidden_size must match num_layers if num_layers is specified.")
 
         params = {} if inplace is None else {"inplace": inplace}
 
@@ -84,10 +81,7 @@ class MLP(torch.nn.Module):
         if isinstance(output_heads, int):
             output_heads = {"output": output_heads}
         self.heads = torch.nn.ModuleDict(
-            {
-                name: torch.nn.Linear(hidden_size[-1], output_size)
-                for name, output_size in output_heads.items()
-            }
+            {name: torch.nn.Linear(hidden_size[-1], output_size) for name, output_size in output_heads.items()}
         )
 
     @classmethod

ezmsg/learn/model/refit_kalman.py

@@ -83,17 +83,12 @@ class RefitKalmanFilter:
         if Y_state.ndim != 2:
             raise ValueError(f"State vector must be 2D, got {Y_state.ndim}D")
 
-        if (
-            not hasattr(self, "H_observation_matrix")
-            or self.H_observation_matrix is None
-        ):
+        if not hasattr(self, "H_observation_matrix") or self.H_observation_matrix is None:
             raise ValueError("Model must be fitted before refitting")
 
         expected_states = self.H_observation_matrix.shape[1]
         if Y_state.shape[1] != expected_states:
-            raise ValueError(
-                f"State vector has {Y_state.shape[1]} dimensions, expected {expected_states}"
-            )
+            raise ValueError(f"State vector has {Y_state.shape[1]} dimensions, expected {expected_states}")
 
     def fit(self, X_train, y_train):
         """
@@ -121,15 +116,11 @@ class RefitKalmanFilter:
         X2 = X[1:, :]  # x_{t+1}
         X1 = X[:-1, :]  # x_t
         A = X2.T @ X1 @ np.linalg.inv(X1.T @ X1)  # Transition matrix
-        W = (
-            (X2 - X1 @ A.T).T @ (X2 - X1 @ A.T) / (n_samples - 1)
-        )  # Covariance of transition matrix
+        W = (X2 - X1 @ A.T).T @ (X2 - X1 @ A.T) / (n_samples - 1)  # Covariance of transition matrix
 
         # Calculate the measurement matrix (from x_t to z_t) using least-squares
         H = Z.T @ X @ np.linalg.inv(X.T @ X)  # Measurement matrix
-        Q = (
-            (Z - X @ H.T).T @ (Z - X @ H.T) / Z.shape[0]
-        )  # Covariance of measurement matrix
+        Q = (Z - X @ H.T).T @ (Z - X @ H.T) / Z.shape[0]  # Covariance of measurement matrix
 
         self.A_state_transition_matrix = A
         self.W_process_noise_covariance = W * self.process_noise_scale
@@ -179,15 +170,11 @@ class RefitKalmanFilter:
         if intention_velocity_indices is None:
             # Assume (x, y, vx, vy)
             vel_idx = 2 if Y_state.shape[1] >= 4 else 0
-            print(
-                f"[RefitKalmanFilter] No velocity index provided — defaulting to {vel_idx}"
-            )
+            print(f"[RefitKalmanFilter] No velocity index provided — defaulting to {vel_idx}")
         else:
             if isinstance(intention_velocity_indices, (list, tuple)):
                 if len(intention_velocity_indices) != 1:
-                    raise ValueError(
-                        "Only one velocity start index should be provided."
-                    )
+                    raise ValueError("Only one velocity start index should be provided.")
                 vel_idx = intention_velocity_indices[0]
             else:
                 vel_idx = intention_velocity_indices
@@ -198,18 +185,14 @@ class RefitKalmanFilter:
         else:
             intended_states = Y_state.copy()
             # Calculate intended velocities for each sample
-            for i, (state, pos, target) in enumerate(
-                zip(Y_state, cursor_positions, target_positions)
-            ):
+            for i, (state, pos, target) in enumerate(zip(Y_state, cursor_positions, target_positions)):
                 is_hold = hold_indices[i] if hold_indices is not None else False
 
                 if is_hold:
                     # During hold periods, intended velocity is zero
                     intended_states[i, vel_idx : vel_idx + 2] = 0.0
                     if i > 0:
-                        intended_states[i, :2] = intended_states[
-                            i - 1, :2
-                        ]  # Same position as previous
+                        intended_states[i, :2] = intended_states[i - 1, :2]  # Same position as previous
                 else:
                     # Calculate direction to target
                     to_target = target - pos
@@ -228,9 +211,7 @@ class RefitKalmanFilter:
                         intended_states[i, vel_idx : vel_idx + 2] = intended_velocity
                     # If target is very close, keep original velocity
                     else:
-                        intended_states[i, vel_idx : vel_idx + 2] = state[
-                            vel_idx : vel_idx + 2
-                        ]
+                        intended_states[i, vel_idx : vel_idx + 2] = state[vel_idx : vel_idx + 2]
 
         intended_states = np.array(intended_states)
         Z = np.array(X_neural)
@@ -258,7 +239,8 @@ class RefitKalmanFilter:
         Raises:
             LinAlgError: If the Riccati equation cannot be solved or matrix operations fail.
         """
-        ## TODO: consider removing non-steady-state for compute_gain() - non_steady_state updates will occur during predict() and update()
+        # TODO: consider removing non-steady-state for compute_gain() -
+        #  non_steady_state updates will occur during predict() and update()
         # if self.steady_state:
         try:
             # Try with original matrices
@@ -272,9 +254,7 @@ class RefitKalmanFilter:
                 self.P_state_covariance
                 @ self.H_observation_matrix.T
                 @ np.linalg.inv(
-                    self.H_observation_matrix
-                    @ self.P_state_covariance
-                    @ self.H_observation_matrix.T
+                    self.H_observation_matrix @ self.P_state_covariance @ self.H_observation_matrix.T
                     + self.Q_measurement_noise_covariance
                 )
             )
@@ -284,9 +264,7 @@ class RefitKalmanFilter:
             # W_reg = self.W_process_noise_covariance + 1e-7 * np.eye(
             # self.W_process_noise_covariance.shape[0]
             # )
-            Q_reg = self.Q_measurement_noise_covariance + 1e-7 * np.eye(
-                self.Q_measurement_noise_covariance.shape[0]
-            )
+            Q_reg = self.Q_measurement_noise_covariance + 1e-7 * np.eye(self.Q_measurement_noise_covariance.shape[0])
 
             try:
                 self.P_state_covariance = solve_discrete_are(
@@ -299,19 +277,14 @@ class RefitKalmanFilter:
                     self.P_state_covariance
                     @ self.H_observation_matrix.T
                     @ np.linalg.inv(
-                        self.H_observation_matrix
-                        @ self.P_state_covariance
-                        @ self.H_observation_matrix.T
-                        + Q_reg
+                        self.H_observation_matrix @ self.P_state_covariance @ self.H_observation_matrix.T + Q_reg
                     )
                 )
                 print("Warning: Used regularized matrices for DARE solution")
             except LinAlgError:
                 # Fallback to identity or manual initialization
                 print("Warning: DARE failed, using identity covariance")
-                self.P_state_covariance = np.eye(
-                    self.A_state_transition_matrix.shape[0]
-                )
+                self.P_state_covariance = np.eye(self.A_state_transition_matrix.shape[0])
 
         # else:
         #     n_states = self.A_state_transition_matrix.shape[0]
@@ -349,9 +322,7 @@ class RefitKalmanFilter:
             return x_predicted, None
         else:
             P_predicted = self.alpha_fading_memory**2 * (
-                self.A_state_transition_matrix
-                @ self.P_state_covariance
-                @ self.A_state_transition_matrix.T
+                self.A_state_transition_matrix @ self.P_state_covariance @ self.A_state_transition_matrix.T
                 + self.W_process_noise_covariance
             )
             return x_predicted, P_predicted
@@ -376,10 +347,7 @@ class RefitKalmanFilter:
 
         # Non-steady-state mode
         # System uncertainty
-        S = (
-            self.H_observation_matrix @ P_predicted @ self.H_observation_matrix.T
-            + self.Q_measurement_noise_covariance
-        )
+        S = self.H_observation_matrix @ P_predicted @ self.H_observation_matrix.T + self.Q_measurement_noise_covariance
 
         # Kalman gain
         K = P_predicted @ self.H_observation_matrix.T @ np.linalg.pinv(S)

ezmsg/learn/nlin_model/mlp.py

@@ -1,6 +1,10 @@
 from ..model.mlp_old import MLP as MLP
+from ..process.mlp_old import (
+    MLPProcessor as MLPProcessor,
+)
 from ..process.mlp_old import (
     MLPSettings as MLPSettings,
+)
+from ..process.mlp_old import (
     MLPState as MLPState,
-    MLPProcessor as MLPProcessor,
 )

ezmsg/learn/process/adaptive_linear_regressor.py

@@ -1,16 +1,15 @@
 from dataclasses import field
 
+import ezmsg.core as ez
 import numpy as np
 import pandas as pd
-import river.optim
 import river.linear_model
+import river.optim
 import sklearn.base
-import ezmsg.core as ez
-from ezmsg.sigproc.sampler import SampleMessage
-from ezmsg.sigproc.base import (
-    processor_state,
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
+    processor_state,
 )
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
@@ -39,9 +38,7 @@ class AdaptiveLinearRegressorTransformer(
 ):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.settings = replace(
-            self.settings, model_type=AdaptiveLinearRegressor(self.settings.model_type)
-        )
+        self.settings = replace(self.settings, model_type=AdaptiveLinearRegressor(self.settings.model_type))
         b_river = self.settings.model_type in [
             AdaptiveLinearRegressor.LINEAR,
             AdaptiveLinearRegressor.LOGISTIC,
@@ -49,9 +46,7 @@ class AdaptiveLinearRegressorTransformer(
         if b_river:
             self.settings.model_kwargs["l2"] = self.settings.model_kwargs.get("l2", 0.0)
             if "learn_rate" in self.settings.model_kwargs:
-                self.settings.model_kwargs["optimizer"] = river.optim.SGD(
-                    self.settings.model_kwargs.pop("learn_rate")
-                )
+                self.settings.model_kwargs["optimizer"] = river.optim.SGD(self.settings.model_kwargs.pop("learn_rate"))
 
         if self.settings.settings_path is not None:
             # Load model from file
@@ -69,9 +64,7 @@ class AdaptiveLinearRegressorTransformer(
                 print("TODO: Override sklearn model with kwargs")
         else:
             # Build model from scratch.
-            regressor_klass = get_regressor(
-                RegressorType.ADAPTIVE, self.settings.model_type
-            )
+            regressor_klass = get_regressor(RegressorType.ADAPTIVE, self.settings.model_type)
             self.state.model = regressor_klass(**self.settings.model_kwargs)
 
     def _hash_message(self, message: AxisArray) -> int:
@@ -84,37 +77,30 @@ class AdaptiveLinearRegressorTransformer(
         #  .template is updated in partial_fit
         pass
 
-    def partial_fit(self, message: SampleMessage) -> None:
-        if np.any(np.isnan(message.sample.data)):
+    def partial_fit(self, message: AxisArray) -> None:
+        if np.any(np.isnan(message.data)):
             return
 
         if self.settings.model_type in [
             AdaptiveLinearRegressor.LINEAR,
             AdaptiveLinearRegressor.LOGISTIC,
         ]:
-            x = pd.DataFrame.from_dict(
-                {
-                    k: v
-                    for k, v in zip(
-                        message.sample.axes["ch"].data, message.sample.data.T
-                    )
-                }
-            )
+            x = pd.DataFrame.from_dict({k: v for k, v in zip(message.axes["ch"].data, message.data.T)})
             y = pd.Series(
-                data=message.trigger.value.data[:, 0],
-                name=message.trigger.value.axes["ch"].data[0],
+                data=message.attrs["trigger"].value.data[:, 0],
+                name=message.attrs["trigger"].value.axes["ch"].data[0],
             )
             self.state.model.learn_many(x, y)
         else:
-            X = message.sample.data
-            if message.sample.get_axis_idx("time") != 0:
-                X = np.moveaxis(X, message.sample.get_axis_idx("time"), 0)
-            self.state.model.partial_fit(X, message.trigger.value.data)
+            X = message.data
+            if message.get_axis_idx("time") != 0:
+                X = np.moveaxis(X, message.get_axis_idx("time"), 0)
+            self.state.model.partial_fit(X, message.attrs["trigger"].value.data)
 
         self.state.template = replace(
-            message.trigger.value,
-            data=np.empty_like(message.trigger.value.data),
-            key=message.trigger.value.key + "_pred",
+            message.attrs["trigger"].value,
+            data=np.empty_like(message.attrs["trigger"].value.data),
+            key=message.attrs["trigger"].value.key + "_pred",
         )
 
     def _process(self, message: AxisArray) -> AxisArray | None:
@@ -127,9 +113,7 @@ class AdaptiveLinearRegressorTransformer(
                 AdaptiveLinearRegressor.LOGISTIC,
             ]:
                 # convert msg_in.data to something appropriate for river
-                x = pd.DataFrame.from_dict(
-                    {k: v for k, v in zip(message.axes["ch"].data, message.data.T)}
-                )
+                x = pd.DataFrame.from_dict({k: v for k, v in zip(message.axes["ch"].data, message.data.T)})
                 preds = self.state.model.predict_many(x).values
             else:
                 preds = self.state.model.predict(message.data)

ezmsg/learn/process/base.py

@@ -1,7 +1,7 @@
 import inspect
 import json
-from pathlib import Path
 import typing
+from pathlib import Path
 
 import ezmsg.core as ez
 import torch
@@ -32,9 +32,7 @@ class ModelInitMixin:
         for key, value in config.items():
             if key in model_kwargs:
                 if model_kwargs[key] != value:
-                    ez.logger.warning(
-                        f"Config parameter {key} ({value}) differs from settings ({model_kwargs[key]})."
-                    )
+                    ez.logger.warning(f"Config parameter {key} ({value}) differs from settings ({model_kwargs[key]}).")
             else:
                 ez.logger.warning(f"Config parameter {key} is not in model_kwargs.")
             model_kwargs[key] = value
@@ -44,7 +42,8 @@ class ModelInitMixin:
         filtered_out = set(kwargs.keys()) - {k for k in valid_params if k != "self"}
         if filtered_out:
             ez.logger.warning(
-                f"Ignoring unexpected model parameters not accepted by {model_class.__name__} constructor: {sorted(filtered_out)}"
+                "Ignoring unexpected model parameters not accepted by"
+                f"{model_class.__name__} constructor: {sorted(filtered_out)}"
             )
         # Keep all valid parameters, including None values, so checkpoint-inferred values can overwrite them
         return {k: v for k, v in kwargs.items() if k in valid_params and k != "self"}
@@ -92,22 +91,16 @@ class ModelInitMixin:
                     config = json.load(f)
                 self._merge_config(model_kwargs, config)
             except Exception as e:
-                raise RuntimeError(
-                    f"Failed to load config from {config_path}: {str(e)}"
-                )
+                raise RuntimeError(f"Failed to load config from {config_path}: {str(e)}")
 
         # If a checkpoint file is provided, load it.
         if checkpoint_path:
             checkpoint_path = Path(checkpoint_path)
             if not checkpoint_path.exists():
                 ez.logger.error(f"Checkpoint path {checkpoint_path} does not exist.")
-                raise FileNotFoundError(
-                    f"Checkpoint path {checkpoint_path} does not exist."
-                )
+                raise FileNotFoundError(f"Checkpoint path {checkpoint_path} does not exist.")
             try:
-                checkpoint = torch.load(
-                    checkpoint_path, map_location=device, weights_only=weights_only
-                )
+                checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=weights_only)
 
                 if "config" in checkpoint:
                     config = checkpoint["config"]
@@ -126,20 +119,14 @@ class ModelInitMixin:
                     "infer_config_from_state_dict",
                     lambda _state_dict: {},  # Default to empty dict if not defined
                 )
-                infer_kwargs = (
-                    {"rnn_type": model_kwargs["rnn_type"]}
-                    if "rnn_type" in model_kwargs
-                    else {}
-                )
+                infer_kwargs = {"rnn_type": model_kwargs["rnn_type"]} if "rnn_type" in model_kwargs else {}
                 self._merge_config(
                     model_kwargs,
                     infer_config(state_dict, **infer_kwargs),
                 )
 
             except Exception as e:
-                raise RuntimeError(
-                    f"Failed to load checkpoint from {checkpoint_path}: {str(e)}"
-                )
+                raise RuntimeError(f"Failed to load checkpoint from {checkpoint_path}: {str(e)}")
 
         # Filter model_kwargs to only include valid parameters for the model class
         filtered_kwargs = self._filter_model_kwargs(model_class, model_kwargs)
@@ -156,18 +143,12 @@ class ModelInitMixin:
             if state_dict_prefix:
                 # If a prefix is provided, filter the state_dict keys
                 state_dict = {
-                    k[len(state_dict_prefix) :]: v
-                    for k, v in state_dict.items()
-                    if k.startswith(state_dict_prefix)
+                    k[len(state_dict_prefix) :]: v for k, v in state_dict.items() if k.startswith(state_dict_prefix)
                 }
             # Load the model weights
-            missing, unexpected = model.load_state_dict(
-                state_dict, strict=False, assign=True
-            )
+            missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
             if missing or unexpected:
-                ez.logger.warning(
-                    f"Partial load: missing keys: {missing}, unexpected keys: {unexpected}"
-                )
+                ez.logger.warning(f"Partial load: missing keys: {missing}, unexpected keys: {unexpected}")
 
         model.to(device)
         return model