Wrappers

Model adapters normalizing predict / predict_proba across frameworks.

modeltest.wrappers

Model adapters.

TestContext talks to models through a small common interface:

predict(X)        -> class labels
predict_proba(X)  -> probability estimates (optional)

wrap(model) returns a normalized object implementing this interface for the given model, dispatching on the framework:

• scikit-learn / numpy (BaseEstimator)
• PyTorch (torch.nn.Module) — lazily imported
• Keras / TensorFlow — lazily imported

If the model can't be recognised, we assume it already follows the common interface (e.g. a user-provided wrapper) and let predict/predict_proba resolve dynamically.

ModelWrapper(model)

Normalized interface all scenarios use to query a model.

Wrap model and expose the normalized interface.

Source code in modeltest/wrappers.py

def __init__(self, model: Any):
    """Wrap ``model`` and expose the normalized interface."""
    self.model = model
    self._predict_fn: Optional[Callable] = None
    self._proba_fn: Optional[Callable] = None

feature_names_in_ property

Expose the model's feature names when available (for filtering).

predict(X)

Return class labels for X. Concrete adapters implement this.

Source code in modeltest/wrappers.py

def predict(self, X: Any) -> np.ndarray:
    """Return class labels for ``X``. Concrete adapters implement this."""
    raise NotImplementedError

predict_proba(X)

Return probability estimates for X, or None when the underlying model does not support them.

Source code in modeltest/wrappers.py

def predict_proba(self, X: Any) -> Optional[np.ndarray]:
    """Return probability estimates for ``X``, or ``None`` when the
    underlying model does not support them."""
    raise NotImplementedError

SklearnModel(model)

Bases: ModelWrapper

Adapter for scikit-learn style estimators exposing predict.

Source code in modeltest/wrappers.py

def __init__(self, model: Any):
    """Wrap ``model`` and expose the normalized interface."""
    self.model = model
    self._predict_fn: Optional[Callable] = None
    self._proba_fn: Optional[Callable] = None

predict(X)

Delegate the prediction to the wrapped estimator.

Source code in modeltest/wrappers.py

def predict(self, X: Any) -> np.ndarray:
    """Delegate the prediction to the wrapped estimator."""
    return self.model.predict(X)

predict_proba(X)

Return the estimator's probabilities when available, else None.

Source code in modeltest/wrappers.py

def predict_proba(self, X: Any) -> Optional[np.ndarray]:
    """Return the estimator's probabilities when available, else
    ``None``."""
    if hasattr(self.model, "predict_proba"):
        return np.asarray(self.model.predict_proba(X))
    return None

SklearnClassifier(model)

Bases: SklearnModel

Adapter for classifiers: labels + probabilities.

Source code in modeltest/wrappers.py

def __init__(self, model: Any):
    """Wrap ``model`` and expose the normalized interface."""
    self.model = model
    self._predict_fn: Optional[Callable] = None
    self._proba_fn: Optional[Callable] = None

predict_proba(X)

Return classifier probabilities (assumed available).

Source code in modeltest/wrappers.py

def predict_proba(self, X: Any) -> Optional[np.ndarray]:
    """Return classifier probabilities (assumed available)."""
    return np.asarray(self.model.predict_proba(X))

TorchModel(model, *, input_key='images', device=None)

Bases: ModelWrapper

Adapter for a PyTorch nn.Module trained for classification.

predict returns argmax class indices; predict_proba returns softmax probabilities. Expects X as a numpy array or torch tensor.

Configure the torch adapter.

Args: model: A torch.nn.Module classifier. input_key: Reserved for models expecting dict inputs. device: Optional device to move inputs to before forward.

Source code in modeltest/wrappers.py

def __init__(self, model: Any, *, input_key: str = "images", device: Any = None):
    """Configure the torch adapter.

    Args:
        model: A ``torch.nn.Module`` classifier.
        input_key: Reserved for models expecting dict inputs.
        device: Optional device to move inputs to before forward.
    """
    super().__init__(model)
    self.input_key = input_key
    self.device = device

predict(X)

Forward X (in inference mode) and return argmax class indices.

Source code in modeltest/wrappers.py

def predict(self, X: Any) -> np.ndarray:
    """Forward ``X`` (in inference mode) and return argmax class
    indices."""
    out = self._forward(X)
    if hasattr(out, "detach"):
        data = out.detach().cpu().numpy()
    else:
        data = np.asarray(out).squeeze()
    return np.asarray(data).argmax(axis=1).astype(int)

predict_proba(X)

Forward X and return softmax-normalized probabilities.

Source code in modeltest/wrappers.py

def predict_proba(self, X: Any) -> Optional[np.ndarray]:
    """Forward ``X`` and return softmax-normalized probabilities."""
    out = self._forward(X)
    data = out.detach().cpu().numpy() if hasattr(out, "detach") else np.asarray(out)
    data = np.asarray(data)
    if data.ndim == 1:
        data = data[:, None]
    exp = np.exp(data - data.max(axis=1, keepdims=True))
    return exp / exp.sum(axis=1, keepdims=True)

KerasModel(model, *, multiclass=False)

Bases: ModelWrapper

Adapter for a compiled Keras/TensorFlow model.

Configure the Keras adapter.

Args: model: A compiled tf.keras.Model. multiclass: Force argmax decoding even when the model has two outputs (default: threshold at 0.5 for single-output).

Source code in modeltest/wrappers.py

def __init__(self, model: Any, *, multiclass: bool = False):
    """Configure the Keras adapter.

    Args:
        model: A compiled ``tf.keras.Model``.
        multiclass: Force argmax decoding even when the model has two
            outputs (default: threshold at 0.5 for single-output).
    """
    super().__init__(model)
    self.multiclass = multiclass

predict(X)

Return class labels: argmax for multiclass outputs, 0.5 threshold for single-output models.

Source code in modeltest/wrappers.py

def predict(self, X: Any) -> np.ndarray:
    """Return class labels: argmax for multiclass outputs, 0.5
    threshold for single-output models."""
    proba = np.asarray(self.model.predict(X, verbose=0))
    if self.multiclass or proba.ndim == 2 and proba.shape[1] > 2:
        return proba.argmax(axis=1)
    return (proba > 0.5).astype(int).ravel()

predict_proba(X)

Return the model's raw output as probabilities.

Source code in modeltest/wrappers.py

def predict_proba(self, X: Any) -> Optional[np.ndarray]:
    """Return the model's raw output as probabilities."""
    return np.asarray(self.model.predict(X, verbose=0))

wrap(model, **kwargs)

Return a normalized :class:ModelWrapper for model.

kwargs are forwarded to the adapter (e.g. input_key for Torch, multiclass for Keras). If the model is already a ModelWrapper it is returned unchanged.

Source code in modeltest/wrappers.py

def wrap(model: Any, **kwargs: Any) -> ModelWrapper:
    """Return a normalized :class:`ModelWrapper` for ``model``.

    ``kwargs`` are forwarded to the adapter (e.g. ``input_key`` for Torch,
    ``multiclass`` for Keras). If the model is already a ``ModelWrapper`` it is
    returned unchanged.
    """
    if isinstance(model, ModelWrapper):
        return model

    # Torch
    try:
        import torch  # noqa: F401

        if isinstance(model, torch.nn.Module):
            return TorchModel(model, **kwargs)
    except ImportError:
        pass

    # Keras / TF
    try:
        import tensorflow as tf  # noqa: F401

        if hasattr(tf, "keras") and isinstance(model, tf.keras.Model):
            return KerasModel(model, **kwargs)
    except ImportError:
        pass

    # scikit-learn / sklearn-style
    try:
        from sklearn.base import BaseEstimator

        if isinstance(model, BaseEstimator):
            if hasattr(model, "predict_proba"):
                return SklearnClassifier(model)
            return SklearnModel(model)
    except ImportError:
        pass

    # Fallback: optimistically assume the common interface.
    return SklearnModel(model)