Spaces:
Running
Running
File size: 5,626 Bytes
e17f3ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """
Base Model Wrapper
==================
Abstract base class for all model wrappers.
Ensures sklearn-compatible interface for consistent evaluation.
Author: UW MSIM Team
Date: November 2025
"""
from abc import ABC, abstractmethod
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
import time
import logging
from typing import Any, Optional
import numpy as np
import pandas as pd
logger = logging.getLogger(__name__)
class BaseModelWrapper(BaseEstimator, ABC):
"""
Base class for all model wrappers.
Ensures sklearn-compatible interface with:
- fit(X, y): Train the model
- predict(X): Make predictions
- predict_proba(X): Predict class probabilities (classification only)
Also tracks timing information:
- fit_time: Time spent in training
- predict_time: Time spent in prediction
Parameters
----------
task_type : str, default='classification'
Type of task: 'classification' or 'regression'
random_state : int, optional
Random seed for reproducibility
"""
def __init__(
self,
task_type: str = 'classification',
random_state: Optional[int] = 42
):
self.task_type = task_type
self.random_state = random_state
self.model = None
self.fit_time: Optional[float] = None
self.predict_time: Optional[float] = None
self.is_fitted: bool = False
@abstractmethod
def fit(self, X: Any, y: Any) -> 'BaseModelWrapper':
"""
Train the model on provided data.
Parameters
----------
X : pd.DataFrame or np.ndarray, shape (n_samples, n_features)
Training features
y : pd.Series or np.ndarray, shape (n_samples,)
Training target
Returns
-------
self : BaseModelWrapper
Returns self for method chaining
"""
pass
@abstractmethod
def predict(self, X: Any) -> np.ndarray:
"""
Make predictions on new data.
Parameters
----------
X : pd.DataFrame or np.ndarray, shape (n_samples, n_features)
Test features
Returns
-------
predictions : np.ndarray, shape (n_samples,)
Predicted values or class labels
"""
pass
def predict_proba(self, X: Any) -> np.ndarray:
"""
Predict class probabilities (classification only).
Parameters
----------
X : pd.DataFrame or np.ndarray, shape (n_samples, n_features)
Test features
Returns
-------
probabilities : np.ndarray, shape (n_samples, n_classes)
Class probabilities
Raises
------
NotImplementedError
If task_type is not 'classification'
ValueError
If model is not fitted
"""
if self.task_type != 'classification':
raise NotImplementedError(
f"predict_proba only available for classification tasks, "
f"got task_type='{self.task_type}'"
)
if not self.is_fitted:
raise ValueError("Model not fitted. Call fit() first.")
start_time = time.time()
proba = self._predict_proba_impl(X)
self.predict_time = time.time() - start_time
return proba
@abstractmethod
def _predict_proba_impl(self, X: Any) -> np.ndarray:
"""
Implementation of predict_proba (model-specific).
Parameters
----------
X : pd.DataFrame or np.ndarray, shape (n_samples, n_features)
Test features
Returns
-------
probabilities : np.ndarray, shape (n_samples, n_classes)
Class probabilities
"""
pass
def get_params(self, deep: bool = True) -> dict:
"""
Get parameters for this estimator (sklearn compatibility).
Parameters
----------
deep : bool, default=True
If True, return parameters for sub-estimators
Returns
-------
params : dict
Parameter names mapped to their values
"""
return {
'task_type': self.task_type,
'random_state': self.random_state
}
def set_params(self, **params) -> 'BaseModelWrapper':
"""
Set parameters for this estimator (sklearn compatibility).
Parameters
----------
**params : dict
Estimator parameters
Returns
-------
self : BaseModelWrapper
Returns self
"""
for key, value in params.items():
setattr(self, key, value)
return self
def _validate_input(self, X: Any, y: Optional[Any] = None):
"""
Validate input data format.
Parameters
----------
X : any
Features
y : any, optional
Target (if provided)
"""
# Convert to pandas if needed
if not isinstance(X, (pd.DataFrame, np.ndarray)):
raise TypeError(
f"X must be pd.DataFrame or np.ndarray, got {type(X)}"
)
if y is not None and not isinstance(y, (pd.Series, np.ndarray)):
raise TypeError(
f"y must be pd.Series or np.ndarray, got {type(y)}"
)
def __repr__(self) -> str:
"""String representation of the model."""
params = self.get_params()
param_str = ', '.join(f"{k}={v}" for k, v in params.items())
return f"{self.__class__.__name__}({param_str})"
|